error -> warning + clean

.github/actions/custom-build-and-push/action.yml

@@ -32,20 +32,16 @@ inputs:
     description: 'Cache destinations'
     required: false
   retry-wait-time:
-    description: 'Time to wait before attempt 2 in seconds'
+    description: 'Time to wait before retry in seconds'
     required: false
-    default: '60'
-  retry-wait-time-2:
-    description: 'Time to wait before attempt 3 in seconds'
-    required: false
-    default: '120'
+    default: '5'
 
 runs:
   using: "composite"
   steps:
-    - name: Build and push Docker image (Attempt 1 of 3)
+    - name: Build and push Docker image (First Attempt)
       id: buildx1
-      uses: docker/build-push-action@v6
+      uses: docker/build-push-action@v5
       continue-on-error: true
       with:
         context: ${{ inputs.context }}
@@ -58,17 +54,16 @@ runs:
         cache-from: ${{ inputs.cache-from }}
         cache-to: ${{ inputs.cache-to }}
 
-    - name: Wait before attempt 2
+    - name: Wait to retry
       if: steps.buildx1.outcome != 'success'
       run: |
         echo "First attempt failed. Waiting ${{ inputs.retry-wait-time }} seconds before retry..."
         sleep ${{ inputs.retry-wait-time }}
       shell: bash
 
-    - name: Build and push Docker image (Attempt 2 of 3)
-      id: buildx2
+    - name: Build and push Docker image (Retry Attempt)
       if: steps.buildx1.outcome != 'success'
-      uses: docker/build-push-action@v6
+      uses: docker/build-push-action@v5
       with:
         context: ${{ inputs.context }}
         file: ${{ inputs.file }}
@@ -79,31 +74,3 @@ runs:
         tags: ${{ inputs.tags }}
         cache-from: ${{ inputs.cache-from }}
         cache-to: ${{ inputs.cache-to }}
-
-    - name: Wait before attempt 3
-      if: steps.buildx1.outcome != 'success' && steps.buildx2.outcome != 'success'
-      run: |
-        echo "Second attempt failed. Waiting ${{ inputs.retry-wait-time-2 }} seconds before retry..."
-        sleep ${{ inputs.retry-wait-time-2 }}
-      shell: bash
-
-    - name: Build and push Docker image (Attempt 3 of 3)
-      id: buildx3
-      if: steps.buildx1.outcome != 'success' && steps.buildx2.outcome != 'success'
-      uses: docker/build-push-action@v6
-      with:
-        context: ${{ inputs.context }}
-        file: ${{ inputs.file }}
-        platforms: ${{ inputs.platforms }}
-        pull: ${{ inputs.pull }}
-        push: ${{ inputs.push }}
-        load: ${{ inputs.load }}
-        tags: ${{ inputs.tags }}
-        cache-from: ${{ inputs.cache-from }}
-        cache-to: ${{ inputs.cache-to }}
-
-    - name: Report failure
-      if: steps.buildx1.outcome != 'success' && steps.buildx2.outcome != 'success' && steps.buildx3.outcome != 'success'
-      run: |
-        echo "All attempts failed. Possible transient infrastucture issues? Try again later or inspect logs for details."
-      shell: bash

.github/pull_request_template.md

@@ -6,24 +6,20 @@
 [Describe the tests you ran to verify your changes]
 
 
-## Accepted Risk (provide if relevant)
-N/A
+## Accepted Risk
+[Any know risks or failure modes to point out to reviewers]
 
 
-## Related Issue(s) (provide if relevant)
-N/A
+## Related Issue(s)
+[If applicable, link to the issue(s) this PR addresses]
 
 
-## Mental Checklist:
-- All of the automated tests pass
-- All PR comments are addressed and marked resolved
-- If there are migrations, they have been rebased to latest main
-- If there are new dependencies, they are added to the requirements
-- If there are new environment variables, they are added to all of the deployment methods
-- If there are new APIs that don't require auth, they are added to PUBLIC_ENDPOINT_SPECS
-- Docker images build and basic functionalities work
-- Author has done a final read through of the PR right before merge
-
-## Backporting (check the box to trigger backport action)
-Note: You have to check that the action passes, otherwise resolve the conflicts manually and tag the patches.
-- [ ] This PR should be backported (make sure to check that the backport attempt succeeds)
+## Checklist:
+- [ ] All of the automated tests pass
+- [ ] All PR comments are addressed and marked resolved
+- [ ] If there are migrations, they have been rebased to latest main
+- [ ] If there are new dependencies, they are added to the requirements
+- [ ] If there are new environment variables, they are added to all of the deployment methods
+- [ ] If there are new APIs that don't require auth, they are added to PUBLIC_ENDPOINT_SPECS
+- [ ] Docker images build and basic functionalities work
+- [ ] Author has done a final read through of the PR right before merge

.github/workflows/docker-build-push-web-container-on-tag.yml

@@ -11,11 +11,8 @@ env:
   
 jobs:
   build:
-    runs-on:
-      - runs-on
-      - runner=${{ matrix.platform == 'linux/amd64' && '8cpu-linux-x64' || '8cpu-linux-arm64' }}
-      - run-id=${{ github.run_id }}
-      - tag=platform-${{ matrix.platform }}
+    runs-on: 
+      group: ${{ matrix.platform == 'linux/amd64' && 'amd64-image-builders' || 'arm64-image-builders' }}
     strategy:
       fail-fast: false
       matrix:

.github/workflows/hotfix-release-branches.yml

@@ -1,172 +0,0 @@
-# This workflow is intended to be manually triggered via the GitHub Action tab.
-# Given a hotfix branch, it will attempt to open a PR to all release branches and
-# by default auto merge them
-
-name: Hotfix release branches
-
-on:
-  workflow_dispatch:
-    inputs:
-      hotfix_commit:
-        description: 'Hotfix commit hash'
-        required: true
-      hotfix_suffix:
-        description: 'Hotfix branch suffix (e.g. hotfix/v0.8-{suffix})'
-        required: true
-      release_branch_pattern:
-        description: 'Release branch pattern (regex)'
-        required: true
-        default: 'release/.*'
-      auto_merge:
-        description: 'Automatically merge the hotfix PRs'
-        required: true
-        type: choice
-        default: 'true'
-        options:
-          - true
-          - false
-          
-jobs:
-  hotfix_release_branches:
-    permissions: write-all
-    # See https://runs-on.com/runners/linux/
-    # use a lower powered instance since this just does i/o to docker hub
-    runs-on: [runs-on,runner=2cpu-linux-x64,"run-id=${{ github.run_id }}"]
-    steps:
-    
-      # needs RKUO_DEPLOY_KEY for write access to merge PR's
-      - name: Checkout Repository
-        uses: actions/checkout@v4
-        with:
-          ssh-key: "${{ secrets.RKUO_DEPLOY_KEY }}"
-          fetch-depth: 0
-          
-      - name: Set up Git user
-        run: |
-          git config user.name "Richard Kuo [bot]"
-          git config user.email "rkuo[bot]@danswer.ai"
-
-      - name: Fetch All Branches
-        run: |
-          git fetch --all --prune
-
-      - name: Verify Hotfix Commit Exists
-        run: |
-          git rev-parse --verify "${{ github.event.inputs.hotfix_commit }}" || { echo "Commit not found: ${{ github.event.inputs.hotfix_commit }}"; exit 1; }
-
-      - name: Get Release Branches
-        id: get_release_branches
-        run: |
-          BRANCHES=$(git branch -r | grep -E "${{ github.event.inputs.release_branch_pattern }}" | sed 's|origin/||' | tr -d ' ')
-          if [ -z "$BRANCHES" ]; then
-            echo "No release branches found matching pattern '${{ github.event.inputs.release_branch_pattern }}'."
-            exit 1
-          fi
-          
-          echo "Found release branches:"
-          echo "$BRANCHES"
-          
-          # Join the branches into a single line separated by commas
-          BRANCHES_JOINED=$(echo "$BRANCHES" | tr '\n' ',' | sed 's/,$//')
-
-          # Set the branches as an output
-          echo "branches=$BRANCHES_JOINED" >> $GITHUB_OUTPUT
-
-      # notes on all the vagaries of wiring up automated PR's
-      # https://github.com/peter-evans/create-pull-request/blob/main/docs/concepts-guidelines.md#triggering-further-workflow-runs
-      # we must use a custom token for GH_TOKEN to trigger the subsequent PR checks
-      - name: Create and Merge Pull Requests to Matching Release Branches
-        env:
-          HOTFIX_COMMIT: ${{ github.event.inputs.hotfix_commit }}
-          HOTFIX_SUFFIX: ${{ github.event.inputs.hotfix_suffix }}
-          AUTO_MERGE: ${{ github.event.inputs.auto_merge }}
-          GH_TOKEN: ${{ secrets.RKUO_PERSONAL_ACCESS_TOKEN }}
-        run: |
-          # Get the branches from the previous step
-          BRANCHES="${{ steps.get_release_branches.outputs.branches }}"
-
-          # Convert BRANCHES to an array
-          IFS=$',' read -ra BRANCH_ARRAY <<< "$BRANCHES"
-
-          # Loop through each release branch and create and merge a PR
-          for RELEASE_BRANCH in "${BRANCH_ARRAY[@]}"; do
-            echo "Processing $RELEASE_BRANCH..."
-            
-            # Parse out the release version by removing "release/" from the branch name
-            RELEASE_VERSION=${RELEASE_BRANCH#release/}
-            echo "Release version parsed: $RELEASE_VERSION"
-            
-            HOTFIX_BRANCH="hotfix/${RELEASE_VERSION}-${HOTFIX_SUFFIX}"
-            echo "Creating PR from $HOTFIX_BRANCH to $RELEASE_BRANCH"
-
-            # Checkout the release branch
-            echo "Checking out $RELEASE_BRANCH"
-            git checkout "$RELEASE_BRANCH"
-
-            # Create the new hotfix branch
-            if git rev-parse --verify "$HOTFIX_BRANCH" >/dev/null 2>&1; then
-              echo "Hotfix branch $HOTFIX_BRANCH already exists. Skipping branch creation."
-            else
-              echo "Branching $RELEASE_BRANCH to $HOTFIX_BRANCH"
-              git checkout -b "$HOTFIX_BRANCH"
-            fi
-            
-            # Check if the hotfix commit is a merge commit
-            if git rev-list --merges -n 1 "$HOTFIX_COMMIT" >/dev/null 2>&1; then
-              # -m 1 uses the target branch as the base (which is what we want)
-              echo "Hotfix commit $HOTFIX_COMMIT is a merge commit, using -m 1 for cherry-pick"
-              CHERRY_PICK_CMD="git cherry-pick -m 1 $HOTFIX_COMMIT"
-            else
-              CHERRY_PICK_CMD="git cherry-pick $HOTFIX_COMMIT"
-            fi
-
-            # Perform the cherry-pick
-            echo "Executing: $CHERRY_PICK_CMD"
-            eval "$CHERRY_PICK_CMD"
-
-            if [ $? -ne 0 ]; then
-              echo "Cherry-pick failed for $HOTFIX_COMMIT on $HOTFIX_BRANCH. Aborting..."
-              git cherry-pick --abort
-              continue
-            fi
-
-            # Push the hotfix branch to the remote
-            echo "Pushing $HOTFIX_BRANCH..."
-            git push origin "$HOTFIX_BRANCH"
-            echo "Hotfix branch $HOTFIX_BRANCH created and pushed."
-            
-            # Check if PR already exists
-            EXISTING_PR=$(gh pr list --head "$HOTFIX_BRANCH" --base "$RELEASE_BRANCH" --state open --json number --jq '.[0].number')
-
-            if [ -n "$EXISTING_PR" ]; then
-              echo "An open PR already exists: #$EXISTING_PR. Skipping..."
-              continue
-            fi
-            
-            # Create a new PR and capture the output
-            PR_OUTPUT=$(gh pr create --title "Merge $HOTFIX_BRANCH into $RELEASE_BRANCH" \
-              --body "Automated PR to merge \`$HOTFIX_BRANCH\` into \`$RELEASE_BRANCH\`." \
-              --head "$HOTFIX_BRANCH" --base "$RELEASE_BRANCH")
-
-            # Extract the URL from the output
-            PR_URL=$(echo "$PR_OUTPUT" | grep -Eo 'https://github.com/[^ ]+')
-            echo "Pull request created: $PR_URL"
-
-            # Extract PR number from URL
-            PR_NUMBER=$(basename "$PR_URL")
-            echo "Pull request created: $PR_NUMBER"
-
-            if [ "$AUTO_MERGE" == "true" ]; then
-              echo "Attempting to merge pull request #$PR_NUMBER"
-
-              # Attempt to merge the PR
-              gh pr merge "$PR_NUMBER" --merge --auto --delete-branch
-
-              if [ $? -eq 0 ]; then
-                echo "Pull request #$PR_NUMBER merged successfully."
-              else
-                # Optionally, handle the error or continue
-                echo "Failed to merge pull request #$PR_NUMBER."
-              fi
-            fi
-          done

.github/workflows/nightly-close-stale-issues.yml

@@ -1,23 +0,0 @@
-name: 'Nightly - Close stale issues and PRs'
-on:
-  schedule:
-    - cron: '0 11 * * *' # Runs every day at 3 AM PST / 4 AM PDT / 11 AM UTC
-
-permissions:
-  # contents: write # only for delete-branch option
-  issues: write
-  pull-requests: write
-  
-jobs:
-  stale:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/stale@v9
-        with:
-          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.'
-          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.'
-          close-issue-message: 'This issue was closed because it has been stalled for 90 days with no activity.'
-          close-pr-message: 'This PR was closed because it has been stalled for 90 days with no activity.'
-          days-before-stale: 75
-#           days-before-close: 90  # uncomment after we test stale behavior
-          

.github/workflows/pr-Integration-tests.yml

@@ -72,7 +72,7 @@ jobs:
           load: true
           cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }}
           cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max
-      
+
       - name: Build integration test Docker image
         uses: ./.github/actions/custom-build-and-push
         with:
@@ -85,58 +85,7 @@ jobs:
           cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }}
           cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max
 
-      # Start containers for multi-tenant tests
-      - name: Start Docker containers for multi-tenant tests
-        run: |
-          cd deployment/docker_compose
-          ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \
-          MULTI_TENANT=true \
-          AUTH_TYPE=basic \
-          REQUIRE_EMAIL_VERIFICATION=false \
-          DISABLE_TELEMETRY=true \
-          IMAGE_TAG=test \
-          docker compose -f docker-compose.dev.yml -p danswer-stack up -d
-        id: start_docker_multi_tenant
-
-      # In practice, `cloud` Auth type would require OAUTH credentials to be set.
-      - name: Run Multi-Tenant Integration Tests
-        run: |
-          echo "Running integration tests..."
-          docker run --rm --network danswer-stack_default \
-            --name test-runner \
-            -e POSTGRES_HOST=relational_db \
-            -e POSTGRES_USER=postgres \
-            -e POSTGRES_PASSWORD=password \
-            -e POSTGRES_DB=postgres \
-            -e VESPA_HOST=index \
-            -e REDIS_HOST=cache \
-            -e API_SERVER_HOST=api_server \
-            -e OPENAI_API_KEY=${OPENAI_API_KEY} \
-            -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
-            -e TEST_WEB_HOSTNAME=test-runner \
-            -e AUTH_TYPE=cloud \
-            -e MULTI_TENANT=true \
-            danswer/danswer-integration:test \
-            /app/tests/integration/multitenant_tests
-        continue-on-error: true
-        id: run_multitenant_tests
-
-      - name: Check multi-tenant test results
-        run: |
-          if [ ${{ steps.run_tests.outcome }} == 'failure' ]; then
-            echo "Integration tests failed. Exiting with error."
-            exit 1
-          else
-            echo "All integration tests passed successfully."
-          fi 
-
-      - name: Stop multi-tenant Docker containers
-        run: |
-          cd deployment/docker_compose
-          docker compose -f docker-compose.dev.yml -p danswer-stack down -v
-
-
-      - name: Start Docker containers 
+      - name: Start Docker containers
         run: |
           cd deployment/docker_compose
           ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \
@@ -181,7 +130,7 @@ jobs:
           done
           echo "Finished waiting for service."
 
-      - name: Run Standard Integration Tests
+      - name: Run integration tests
         run: |
           echo "Running integration tests..."
           docker run --rm --network danswer-stack_default \
@@ -196,8 +145,7 @@ jobs:
             -e OPENAI_API_KEY=${OPENAI_API_KEY} \
             -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
             -e TEST_WEB_HOSTNAME=test-runner \
-            danswer/danswer-integration:test \
-            /app/tests/integration/tests
+            danswer/danswer-integration:test
         continue-on-error: true
         id: run_tests
 
@@ -210,11 +158,6 @@ jobs:
             echo "All integration tests passed successfully."
           fi
 
-      - name: Stop Docker containers
-        run: |
-          cd deployment/docker_compose
-          docker compose -f docker-compose.dev.yml -p danswer-stack down -v
-
       - name: Save Docker logs
         if: success() || failure()
         run: |
@@ -224,7 +167,7 @@ jobs:
       
       - name: Upload logs
         if: success() || failure()
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@v3
         with:
           name: docker-logs
           path: ${{ github.workspace }}/docker-compose.log

.github/workflows/pr-backport-autotrigger.yml

@@ -1,98 +0,0 @@
-name: Backport on Merge
-
-on:
-  pull_request:
-    types: [closed]  # Later we check for merge so only PRs that go in can get backported
-
-permissions:
-  contents: write
-  actions: write  # Allows the workflow to trigger other workflows on push
-
-jobs:
-  backport:
-    if: github.event.pull_request.merged == true
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v3
-        with:
-          fetch-depth: 0  # Fetch all history for all branches and tags
-
-      - name: Set up Git
-        run: |
-          git config --global user.name "yuhongsun96"
-          git config --global user.email "yuhongsun96@gmail.com"
-          # Configure Git to use the PAT for authentication
-          git remote set-url origin https://yuhongsun96:${{ secrets.YUHONG_GH_ACTIONS }}@github.com/${{ github.repository }}.git
-
-      - name: Check for Backport Checkbox
-        id: checkbox-check
-        run: |
-          PR_BODY="${{ github.event.pull_request.body }}"
-          if [[ "$PR_BODY" == *"[x] This PR should be backported"* ]]; then
-            echo "backport=true" >> $GITHUB_OUTPUT
-          else
-            echo "backport=false" >> $GITHUB_OUTPUT
-          fi
-
-      - name: List and sort release branches
-        id: list-branches
-        run: |
-          git fetch --all --tags
-          BRANCHES=$(git for-each-ref --format='%(refname:short)' refs/remotes/origin/release/* | sed 's|origin/release/||' | sort -Vr)
-          BETA=$(echo "$BRANCHES" | head -n 1)
-          STABLE=$(echo "$BRANCHES" | head -n 2 | tail -n 1)
-          echo "beta=release/$BETA" >> $GITHUB_OUTPUT
-          echo "stable=release/$STABLE" >> $GITHUB_OUTPUT
-          # Fetch latest tags for beta and stable
-          LATEST_BETA_TAG=$(git tag -l "v[0-9]*.[0-9]*.[0-9]*-beta.[0-9]*" | grep -E "^v[0-9]+\.[0-9]+\.[0-9]+-beta\.[0-9]+$" | grep -v -- "-cloud" | sort -Vr | head -n 1)
-          LATEST_STABLE_TAG=$(git tag -l "v[0-9]*.[0-9]*.[0-9]*" | grep -E "^v[0-9]+\.[0-9]+\.[0-9]+$" | sort -Vr | head -n 1)
-          # Increment latest beta tag
-          NEW_BETA_TAG=$(echo $LATEST_BETA_TAG | awk -F '[.-]' '{print $1 "." $2 "." $3 "-beta." ($NF+1)}')
-          # Increment latest stable tag
-          NEW_STABLE_TAG=$(echo $LATEST_STABLE_TAG | awk -F '.' '{print $1 "." $2 "." ($3+1)}')
-          echo "latest_beta_tag=$LATEST_BETA_TAG" >> $GITHUB_OUTPUT
-          echo "latest_stable_tag=$LATEST_STABLE_TAG" >> $GITHUB_OUTPUT
-          echo "new_beta_tag=$NEW_BETA_TAG" >> $GITHUB_OUTPUT
-          echo "new_stable_tag=$NEW_STABLE_TAG" >> $GITHUB_OUTPUT
-
-      - name: Echo branch and tag information
-        run: |
-          echo "Beta branch: ${{ steps.list-branches.outputs.beta }}"
-          echo "Stable branch: ${{ steps.list-branches.outputs.stable }}"
-          echo "Latest beta tag: ${{ steps.list-branches.outputs.latest_beta_tag }}"
-          echo "Latest stable tag: ${{ steps.list-branches.outputs.latest_stable_tag }}"
-          echo "New beta tag: ${{ steps.list-branches.outputs.new_beta_tag }}"
-          echo "New stable tag: ${{ steps.list-branches.outputs.new_stable_tag }}"
-
-      - name: Trigger Backport
-        if: steps.checkbox-check.outputs.backport == 'true'
-        run: |
-          set -e
-          echo "Backporting to beta ${{ steps.list-branches.outputs.beta }} and stable ${{ steps.list-branches.outputs.stable }}"
-          # Fetch all history for all branches and tags
-          git fetch --prune
-          # Checkout the beta branch
-          git checkout ${{ steps.list-branches.outputs.beta }}
-          # Cherry-pick the merge commit from the merged PR
-          git cherry-pick -m 1 ${{ github.event.pull_request.merge_commit_sha }} || {
-            echo "Cherry-pick to beta failed due to conflicts."
-            exit 1
-          }
-          # Create new beta tag
-          git tag ${{ steps.list-branches.outputs.new_beta_tag }}
-          # Push the changes and tag to the beta branch
-          git push origin ${{ steps.list-branches.outputs.beta }} --force
-          git push origin ${{ steps.list-branches.outputs.new_beta_tag }}
-          # Checkout the stable branch
-          git checkout ${{ steps.list-branches.outputs.stable }}
-          # Cherry-pick the merge commit from the merged PR
-          git cherry-pick -m 1 ${{ github.event.pull_request.merge_commit_sha }} || {
-            echo "Cherry-pick to stable failed due to conflicts."
-            exit 1
-          }
-          # Create new stable tag
-          git tag ${{ steps.list-branches.outputs.new_stable_tag }}
-          # Push the changes and tag to the stable branch
-          git push origin ${{ steps.list-branches.outputs.stable }} --force
-          git push origin ${{ steps.list-branches.outputs.new_stable_tag }}

.github/workflows/pr-python-checks.yml

@@ -14,10 +14,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v4
+      uses: actions/checkout@v3
 
     - name: Set up Python
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v4
       with:
         python-version: '3.11'
         cache: 'pip'

.github/workflows/pr-python-connector-tests.yml

@@ -32,7 +32,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v4
         with:
           python-version: "3.11"
           cache: "pip"

.github/workflows/pr-python-model-tests.yml

@@ -27,7 +27,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v4
         with:
           python-version: "3.11"
           cache: "pip"

.github/workflows/pr-python-tests.yml

@@ -21,7 +21,7 @@ jobs:
       uses: actions/checkout@v4
 
     - name: Set up Python
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v4
       with:
         python-version: '3.11'
         cache: 'pip'

.github/workflows/pr-quality-checks.yml

@@ -18,6 +18,6 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: "3.11"
-      - uses: pre-commit/action@v3.0.1
+      - uses: pre-commit/action@v3.0.0
         with:
           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) || '' }}

.vscode/launch.template.jsonc

@@ -6,69 +6,19 @@
     // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
     "version": "0.2.0",
     "compounds": [
-		{
-			// Dummy entry used to label the group
-			"name": "--- Compound ---",
-		  	"configurations": [
-		  		"--- Individual ---"
-		  	],
-			"presentation": {
-				 "group": "1",
-			 }
-		},
         {
             "name": "Run All Danswer Services",
             "configurations": [
                 "Web Server",
                 "Model Server",
                 "API Server",
-                "Slack Bot",
-		  		"Celery primary", 
-		  		"Celery light", 
-		  		"Celery heavy", 
-		  		"Celery indexing", 
-		  		"Celery beat",
-            ],
-			"presentation": {
-				 "group": "1",
-			 }
-        },
-		{
-			"name": "Web / Model / API",
-		  	"configurations": [
-		  		"Web Server", 
-		  		"Model Server", 
-		  		"API Server",
-		  	],
-			"presentation": {
-				 "group": "1",
-			 }
-		},
-		{
-			"name": "Celery (all)",
-		  	"configurations": [
-		  		"Celery primary", 
-		  		"Celery light", 
-		  		"Celery heavy", 
-		  		"Celery indexing", 
-		  		"Celery beat"
-		  	],
-			"presentation": {
-				 "group": "1",
-			 }
-		}
+                "Indexing",
+                "Background Jobs",
+                "Slack Bot"
+            ]
+        }
     ],
     "configurations": [
-		{
-	    	// Dummy entry used to label the group
-			"name": "--- Individual ---",
-			"type": "node",
-			"request": "launch",
-			"presentation": {
-				"group": "2",
-				"order": 0
-			}
-		},
         {
             "name": "Web Server",
             "type": "node",
@@ -79,11 +29,7 @@
             "runtimeArgs": [
                 "run", "dev"
             ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "console": "integratedTerminal",
-            "consoleTitle": "Web Server Console"
+            "console": "integratedTerminal"
         },
         {
             "name": "Model Server",
@@ -102,11 +48,7 @@
                 "--reload",
                 "--port",
                 "9000"
-            ],
-            "presentation": {
-				 "group": "2",
-			},
-            "consoleTitle": "Model Server Console"
+            ]
         },
         {
             "name": "API Server",
@@ -126,13 +68,43 @@
                 "--reload",
                 "--port",
                 "8080"
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "API Server Console"
+            ]
         },
-        // For the listener to access the Slack API,
+        {
+            "name": "Indexing",
+            "consoleName": "Indexing",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "danswer/background/update.py",
+            "cwd": "${workspaceFolder}/backend",
+            "envFile": "${workspaceFolder}/.vscode/.env",
+            "env": {
+                "ENABLE_MULTIPASS_INDEXING": "false",
+                "LOG_LEVEL": "DEBUG",
+                "PYTHONUNBUFFERED": "1",
+                "PYTHONPATH": "."
+            }
+        },
+        // Celery and all async jobs, usually would include indexing as well but this is handled separately above for dev
+        {
+            "name": "Background Jobs",
+            "consoleName": "Background Jobs",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "scripts/dev_run_background_jobs.py",
+            "cwd": "${workspaceFolder}/backend",
+            "envFile": "${workspaceFolder}/.vscode/.env",
+            "env": {
+                "LOG_DANSWER_MODEL_INTERACTIONS": "True",
+                "LOG_LEVEL": "DEBUG",
+                "PYTHONUNBUFFERED": "1",
+                "PYTHONPATH": "."
+            },
+            "args": [
+                "--no-indexing"
+            ]
+        },
+        // For the listner to access the Slack API,
         // DANSWER_BOT_SLACK_APP_TOKEN & DANSWER_BOT_SLACK_BOT_TOKEN need to be set in .env file located in the root of the project
         {
             "name": "Slack Bot",
@@ -146,151 +118,7 @@
                 "LOG_LEVEL": "DEBUG",
                 "PYTHONUNBUFFERED": "1",
                 "PYTHONPATH": "."
-            },
-            "presentation": {
-				 "group": "2",
-			},
-            "consoleTitle": "Slack Bot Console"
-        },
-        {
-            "name": "Celery primary",
-            "type": "debugpy",
-            "request": "launch",
-            "module": "celery",
-            "cwd": "${workspaceFolder}/backend",
-            "envFile": "${workspaceFolder}/.vscode/.env",
-            "env": {
-                "LOG_LEVEL": "INFO",
-                "PYTHONUNBUFFERED": "1",
-                "PYTHONPATH": "."
-            },
-            "args": [
-                "-A",
-                "danswer.background.celery.versioned_apps.primary",
-                "worker",
-                "--pool=threads",
-                "--concurrency=4",
-                "--prefetch-multiplier=1",
-                "--loglevel=INFO",
-                "--hostname=primary@%n",
-                "-Q",
-                "celery",
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "Celery primary Console"
-        },
-        {
-            "name": "Celery light",
-            "type": "debugpy",
-            "request": "launch",
-            "module": "celery",
-            "cwd": "${workspaceFolder}/backend",
-            "envFile": "${workspaceFolder}/.vscode/.env",
-            "env": {
-                "LOG_LEVEL": "INFO",
-                "PYTHONUNBUFFERED": "1",
-                "PYTHONPATH": "."
-            },
-            "args": [
-                "-A",
-                "danswer.background.celery.versioned_apps.light",
-                "worker",
-                "--pool=threads",
-                "--concurrency=64",
-                "--prefetch-multiplier=8",
-                "--loglevel=INFO",
-                "--hostname=light@%n",
-                "-Q",
-                "vespa_metadata_sync,connector_deletion",
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "Celery light Console"
-        },
-        {
-            "name": "Celery heavy",
-            "type": "debugpy",
-            "request": "launch",
-            "module": "celery",
-            "cwd": "${workspaceFolder}/backend",
-            "envFile": "${workspaceFolder}/.vscode/.env",
-            "env": {
-                "LOG_LEVEL": "INFO",
-                "PYTHONUNBUFFERED": "1",
-                "PYTHONPATH": "."
-            },
-            "args": [
-                "-A",
-                "danswer.background.celery.versioned_apps.heavy",
-                "worker",
-                "--pool=threads",
-                "--concurrency=4",
-                "--prefetch-multiplier=1",
-                "--loglevel=INFO",
-                "--hostname=heavy@%n",
-                "-Q",
-                "connector_pruning",
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "Celery heavy Console"
-        },
-        {
-            "name": "Celery indexing",
-            "type": "debugpy",
-            "request": "launch",
-            "module": "celery",
-            "cwd": "${workspaceFolder}/backend",
-            "envFile": "${workspaceFolder}/.vscode/.env",
-            "env": {
-                "ENABLE_MULTIPASS_INDEXING": "false",
-                "LOG_LEVEL": "DEBUG",
-                "PYTHONUNBUFFERED": "1",
-                "PYTHONPATH": "."
-            },
-            "args": [
-                "-A",
-                "danswer.background.celery.versioned_apps.indexing",
-                "worker",
-                "--pool=threads",
-                "--concurrency=1",
-                "--prefetch-multiplier=1",
-                "--loglevel=INFO",
-                "--hostname=indexing@%n",
-                "-Q",
-                "connector_indexing",
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "Celery indexing Console"
-        },
-        {
-            "name": "Celery beat",
-            "type": "debugpy",
-            "request": "launch",
-            "module": "celery",
-            "cwd": "${workspaceFolder}/backend",
-            "envFile": "${workspaceFolder}/.vscode/.env",
-            "env": {
-                "LOG_LEVEL": "DEBUG",
-                "PYTHONUNBUFFERED": "1",
-                "PYTHONPATH": "."
-            },
-            "args": [
-                "-A",
-                "danswer.background.celery.versioned_apps.beat",
-                "beat",
-                "--loglevel=INFO",
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "Celery beat Console"
+            }
         },
         {
             "name": "Pytest",
@@ -309,22 +137,8 @@
                 "-v"
                 // Specify a sepcific module/test to run or provide nothing to run all tests
                 //"tests/unit/danswer/llm/answering/test_prune_and_merge.py"
-            ],
-            "presentation": {
-				 "group": "2",
-			 },
-            "consoleTitle": "Pytest Console"
+            ]
         },
-		{
-	    	// Dummy entry used to label the group
-			"name": "--- Tasks ---",
-			"type": "node",
-			"request": "launch",
-			"presentation": {
-				"group": "3",
-				"order": 0
-			}
-		},
         {
             "name": "Clear and Restart External Volumes and Containers",
             "type": "node",
@@ -333,27 +147,7 @@
             "runtimeArgs": ["${workspaceFolder}/backend/scripts/restart_containers.sh"],
             "cwd": "${workspaceFolder}",
             "console": "integratedTerminal",
-            "stopOnEntry": true,
-            "presentation": {
-				 "group": "3",
-			 },
-        },
-        {
-	        // Celery jobs launched through a single background script (legacy)
-	        // Recommend using the "Celery (all)" compound launch instead.
-            "name": "Background Jobs",
-            "consoleName": "Background Jobs",
-            "type": "debugpy",
-            "request": "launch",
-            "program": "scripts/dev_run_background_jobs.py",
-            "cwd": "${workspaceFolder}/backend",
-            "envFile": "${workspaceFolder}/.vscode/.env",
-            "env": {
-                "LOG_DANSWER_MODEL_INTERACTIONS": "True",
-                "LOG_LEVEL": "DEBUG",
-                "PYTHONUNBUFFERED": "1",
-                "PYTHONPATH": "."
-            },
-        },
+            "stopOnEntry": true
+        }
     ]
 }

README.md

@@ -68,13 +68,13 @@ We also have built-in support for deployment on Kubernetes. Files for that can b
 
 ## 🚧 Roadmap
 * Chat/Prompt sharing with specific teammates and user groups.
-* Multimodal model support, chat with images, video etc.
+* Multi-Model model support, chat with images, video etc.
 * Choosing between LLMs and parameters during chat session.
 * Tool calling and agent configurations options.
 * Organizational understanding and ability to locate and suggest experts from your team.
 
 
-## Other Notable Benefits of Danswer
+## Other Noteable Benefits of Danswer
 * User Authentication with document level access management.
 * Best in class Hybrid Search across all sources (BM-25 + prefix aware embedding models).
 * Admin Dashboard to configure connectors, document-sets, access, etc.

backend/Dockerfile

@@ -8,12 +8,10 @@ Edition features outside of personal development or testing purposes. Please rea
 founders@danswer.ai for more information. Please visit https://github.com/danswer-ai/danswer"
 
 # Default DANSWER_VERSION, typically overriden during builds by GitHub Actions.
-ARG DANSWER_VERSION=0.8-dev
+ARG DANSWER_VERSION=0.3-dev
 ENV DANSWER_VERSION=${DANSWER_VERSION} \
     DANSWER_RUNNING_IN_DOCKER="true"
 
-ARG CA_CERT_CONTENT=""
-
 RUN echo "DANSWER_VERSION: ${DANSWER_VERSION}"
 # Install system dependencies
 # cmake needed for psycopg (postgres)
@@ -38,17 +36,6 @@ RUN apt-get update && \
     rm -rf /var/lib/apt/lists/* && \
     apt-get clean
 
-
-# Conditionally write the CA certificate and update certificates
-RUN if [ -n "$CA_CERT_CONTENT" ]; then \
-    echo "Adding custom CA certificate"; \
-    echo "$CA_CERT_CONTENT" > /usr/local/share/ca-certificates/my-ca.crt && \
-    chmod 644 /usr/local/share/ca-certificates/my-ca.crt && \
-    update-ca-certificates; \
-else \
-    echo "No custom CA certificate provided"; \
-fi
-
 # Install Python dependencies
 # Remove py which is pulled in by retry, py is not needed and is a CVE
 COPY ./requirements/default.txt /tmp/requirements.txt
@@ -105,7 +92,6 @@ COPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf
 COPY ./danswer /app/danswer
 COPY ./shared_configs /app/shared_configs
 COPY ./alembic /app/alembic
-COPY ./alembic_tenants /app/alembic_tenants
 COPY ./alembic.ini /app/alembic.ini
 COPY supervisord.conf /usr/etc/supervisord.conf
 
@@ -115,7 +101,7 @@ COPY ./scripts/force_delete_connector_by_id.py /app/scripts/force_delete_connect
 # Put logo in assets
 COPY ./assets /app/assets
 
-ENV PYTHONPATH=/app
+ENV PYTHONPATH /app
 
 # Default command which does nothing
 # This container is used by api server and background which specify their own CMD

backend/Dockerfile.model_server

@@ -7,7 +7,7 @@ You can find it at https://hub.docker.com/r/danswer/danswer-model-server. For mo
 visit https://github.com/danswer-ai/danswer."
 
 # Default DANSWER_VERSION, typically overriden during builds by GitHub Actions.
-ARG DANSWER_VERSION=0.8-dev
+ARG DANSWER_VERSION=0.3-dev
 ENV DANSWER_VERSION=${DANSWER_VERSION} \
     DANSWER_RUNNING_IN_DOCKER="true"
 
@@ -55,6 +55,6 @@ COPY ./shared_configs /app/shared_configs
 # Model Server main code
 COPY ./model_server /app/model_server
 
-ENV PYTHONPATH=/app
+ENV PYTHONPATH /app
 
 CMD ["uvicorn", "model_server.main:app", "--host", "0.0.0.0", "--port", "9000"]

backend/alembic.ini

@@ -1,6 +1,6 @@
 # A generic, single database configuration.
 
-[DEFAULT]
+[alembic]
 # path to migration scripts
 script_location = alembic
 
@@ -47,8 +47,7 @@ prepend_sys_path = .
 # version_path_separator = :
 # version_path_separator = ;
 # version_path_separator = space
-version_path_separator = os  
-# Use os.pathsep. Default configuration used for new projects.
+version_path_separator = os  # Use os.pathsep. Default configuration used for new projects.
 
 # set to 'true' to search source files recursively
 # in each "version_locations" directory
@@ -107,12 +106,3 @@ formatter = generic
 [formatter_generic]
 format = %(levelname)-5.5s [%(name)s] %(message)s
 datefmt = %H:%M:%S
-
-
-[alembic]
-script_location = alembic
-version_locations = %(script_location)s/versions
-
-[schema_private]
-script_location = alembic_tenants
-version_locations = %(script_location)s/versions

backend/alembic/env.py

@@ -1,104 +1,104 @@
-from sqlalchemy.engine.base import Connection
-from typing import Any
 import asyncio
 from logging.config import fileConfig
-import logging
 
 from alembic import context
-from sqlalchemy import pool
-from sqlalchemy.ext.asyncio import create_async_engine
-from sqlalchemy.sql import text
-
-from danswer.configs.app_configs import MULTI_TENANT
 from danswer.db.engine import build_connection_string
 from danswer.db.models import Base
+from sqlalchemy import pool
+from sqlalchemy.engine import Connection
+from sqlalchemy.ext.asyncio import create_async_engine
 from celery.backends.database.session import ResultModelBase  # type: ignore
-from danswer.db.engine import get_all_tenant_ids
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from sqlalchemy.schema import SchemaItem
+from sqlalchemy.sql import text
 
 # Alembic Config object
 config = context.config
 
 # Interpret the config file for Python logging.
+# This line sets up loggers basically.
 if config.config_file_name is not None and config.attributes.get(
     "configure_logger", True
 ):
     fileConfig(config.config_file_name)
 
-# Add your model's MetaData object here for 'autogenerate' support
+# Add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
 target_metadata = [Base.metadata, ResultModelBase.metadata]
 
-EXCLUDE_TABLES = {"kombu_queue", "kombu_message"}
 
-# Set up logging
-logger = logging.getLogger(__name__)
-
-
-def include_object(
-    object: Any, name: str, type_: str, reflected: bool, compare_to: Any
-) -> bool:
-    """
-    Determines whether a database object should be included in migrations.
-    Excludes specified tables from migrations.
-    """
-    if type_ == "table" and name in EXCLUDE_TABLES:
-        return False
-    return True
-
-
-def get_schema_options() -> tuple[str, bool, bool]:
-    """
-    Parses command-line options passed via '-x' in Alembic commands.
-    Recognizes 'schema', 'create_schema', and 'upgrade_all_tenants' options.
-    """
+def get_schema_options() -> tuple[str, bool]:
     x_args_raw = context.get_x_argument()
     x_args = {}
     for arg in x_args_raw:
         for pair in arg.split(","):
             if "=" in pair:
                 key, value = pair.split("=", 1)
-                x_args[key.strip()] = value.strip()
-    schema_name = x_args.get("schema", POSTGRES_DEFAULT_SCHEMA)
+                x_args[key] = value
+
+    schema_name = x_args.get("schema", "public")
     create_schema = x_args.get("create_schema", "true").lower() == "true"
-    upgrade_all_tenants = x_args.get("upgrade_all_tenants", "false").lower() == "true"
-
-    if (
-        MULTI_TENANT
-        and schema_name == POSTGRES_DEFAULT_SCHEMA
-        and not upgrade_all_tenants
-    ):
-        raise ValueError(
-            "Cannot run default migrations in public schema when multi-tenancy is enabled. "
-            "Please specify a tenant-specific schema."
-        )
-
-    return schema_name, create_schema, upgrade_all_tenants
+    return schema_name, create_schema
 
 
-def do_run_migrations(
-    connection: Connection, schema_name: str, create_schema: bool
-) -> None:
+EXCLUDE_TABLES = {"kombu_queue", "kombu_message"}
+
+
+def include_object(
+    object: SchemaItem,
+    name: str,
+    type_: str,
+    reflected: bool,
+    compare_to: SchemaItem | None,
+) -> bool:
+    if type_ == "table" and name in EXCLUDE_TABLES:
+        return False
+    return True
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+    Calls to context.execute() here emit the given string to the
+    script output.
     """
-    Executes migrations in the specified schema.
-    """
-    logger.info(f"About to migrate schema: {schema_name}")
+    url = build_connection_string()
+    schema, _ = get_schema_options()
 
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,  # type: ignore
+        literal_binds=True,
+        include_object=include_object,
+        dialect_opts={"paramstyle": "named"},
+        version_table_schema=schema,
+        include_schemas=True,
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def do_run_migrations(connection: Connection) -> None:
+    schema, create_schema = get_schema_options()
     if create_schema:
-        connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"'))
+        connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema}"'))
         connection.execute(text("COMMIT"))
 
-    # Set search_path to the target schema
-    connection.execute(text(f'SET search_path TO "{schema_name}"'))
+    connection.execute(text(f'SET search_path TO "{schema}"'))
 
     context.configure(
         connection=connection,
         target_metadata=target_metadata,  # type: ignore
-        include_object=include_object,
-        version_table_schema=schema_name,
+        version_table_schema=schema,
         include_schemas=True,
         compare_type=True,
         compare_server_default=True,
-        script_location=config.get_main_option("script_location"),
     )
 
     with context.begin_transaction():
@@ -106,98 +106,20 @@ def do_run_migrations(
 
 
 async def run_async_migrations() -> None:
-    """
-    Determines whether to run migrations for a single schema or all schemas,
-    and executes migrations accordingly.
-    """
-    schema_name, create_schema, upgrade_all_tenants = get_schema_options()
-
-    engine = create_async_engine(
+    """Run migrations in 'online' mode."""
+    connectable = create_async_engine(
         build_connection_string(),
         poolclass=pool.NullPool,
     )
 
-    if upgrade_all_tenants:
-        # Run migrations for all tenant schemas sequentially
-        tenant_schemas = get_all_tenant_ids()
+    async with connectable.connect() as connection:
+        await connection.run_sync(do_run_migrations)
 
-        for schema in tenant_schemas:
-            try:
-                logger.info(f"Migrating schema: {schema}")
-                async with engine.connect() as connection:
-                    await connection.run_sync(
-                        do_run_migrations,
-                        schema_name=schema,
-                        create_schema=create_schema,
-                    )
-            except Exception as e:
-                logger.error(f"Error migrating schema {schema}: {e}")
-                raise
-    else:
-        try:
-            logger.info(f"Migrating schema: {schema_name}")
-            async with engine.connect() as connection:
-                await connection.run_sync(
-                    do_run_migrations,
-                    schema_name=schema_name,
-                    create_schema=create_schema,
-                )
-        except Exception as e:
-            logger.error(f"Error migrating schema {schema_name}: {e}")
-            raise
-
-    await engine.dispose()
-
-
-def run_migrations_offline() -> None:
-    """
-    Run migrations in 'offline' mode.
-    """
-    schema_name, _, upgrade_all_tenants = get_schema_options()
-    url = build_connection_string()
-
-    if upgrade_all_tenants:
-        # Run offline migrations for all tenant schemas
-        engine = create_async_engine(url)
-        tenant_schemas = get_all_tenant_ids()
-        engine.sync_engine.dispose()
-
-        for schema in tenant_schemas:
-            logger.info(f"Migrating schema: {schema}")
-            context.configure(
-                url=url,
-                target_metadata=target_metadata,  # type: ignore
-                literal_binds=True,
-                include_object=include_object,
-                version_table_schema=schema,
-                include_schemas=True,
-                script_location=config.get_main_option("script_location"),
-                dialect_opts={"paramstyle": "named"},
-            )
-
-            with context.begin_transaction():
-                context.run_migrations()
-    else:
-        logger.info(f"Migrating schema: {schema_name}")
-        context.configure(
-            url=url,
-            target_metadata=target_metadata,  # type: ignore
-            literal_binds=True,
-            include_object=include_object,
-            version_table_schema=schema_name,
-            include_schemas=True,
-            script_location=config.get_main_option("script_location"),
-            dialect_opts={"paramstyle": "named"},
-        )
-
-        with context.begin_transaction():
-            context.run_migrations()
+    await connectable.dispose()
 
 
 def run_migrations_online() -> None:
-    """
-    Runs migrations in 'online' mode using an asynchronous engine.
-    """
+    """Run migrations in 'online' mode."""
     asyncio.run(run_async_migrations())
 
 

backend/alembic/versions/1b10e1fda030_add_additional_data_to_notifications.py

@@ -1,26 +0,0 @@
-"""add additional data to notifications
-
-Revision ID: 1b10e1fda030
-Revises: 6756efa39ada
-Create Date: 2024-10-15 19:26:44.071259
-
-"""
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-# revision identifiers, used by Alembic.
-revision = "1b10e1fda030"
-down_revision = "6756efa39ada"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "notification", sa.Column("additional_data", postgresql.JSONB(), nullable=True)
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("notification", "additional_data")

backend/alembic/versions/5d12a446f5c0_add_api_version_and_deployment_name_to_.py

@@ -1,30 +0,0 @@
-"""add api_version and deployment_name to search settings
-
-Revision ID: 5d12a446f5c0
-Revises: e4334d5b33ba
-Create Date: 2024-10-08 15:56:07.975636
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "5d12a446f5c0"
-down_revision = "e4334d5b33ba"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "embedding_provider", sa.Column("api_version", sa.String(), nullable=True)
-    )
-    op.add_column(
-        "embedding_provider", sa.Column("deployment_name", sa.String(), nullable=True)
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("embedding_provider", "deployment_name")
-    op.drop_column("embedding_provider", "api_version")

backend/alembic/versions/6756efa39ada_id_uuid_for_chat_session.py

@@ -1,153 +0,0 @@
-"""Migrate chat_session and chat_message tables to use UUID primary keys
-
-Revision ID: 6756efa39ada
-Revises: 5d12a446f5c0
-Create Date: 2024-10-15 17:47:44.108537
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-revision = "6756efa39ada"
-down_revision = "5d12a446f5c0"
-branch_labels = None
-depends_on = None
-
-"""
-This script:
-1. Adds UUID columns to chat_session and chat_message
-2. Populates new columns with UUIDs
-3. Updates foreign key relationships
-4. Removes old integer ID columns
-
-Note: Downgrade will assign new integer IDs, not restore original ones.
-"""
-
-
-def upgrade() -> None:
-    op.execute("CREATE EXTENSION IF NOT EXISTS pgcrypto;")
-
-    op.add_column(
-        "chat_session",
-        sa.Column(
-            "new_id",
-            sa.UUID(as_uuid=True),
-            server_default=sa.text("gen_random_uuid()"),
-            nullable=False,
-        ),
-    )
-
-    op.execute("UPDATE chat_session SET new_id = gen_random_uuid();")
-
-    op.add_column(
-        "chat_message",
-        sa.Column("new_chat_session_id", sa.UUID(as_uuid=True), nullable=True),
-    )
-
-    op.execute(
-        """
-        UPDATE chat_message
-        SET new_chat_session_id = cs.new_id
-        FROM chat_session cs
-        WHERE chat_message.chat_session_id = cs.id;
-        """
-    )
-
-    op.drop_constraint(
-        "chat_message_chat_session_id_fkey", "chat_message", type_="foreignkey"
-    )
-
-    op.drop_column("chat_message", "chat_session_id")
-    op.alter_column(
-        "chat_message", "new_chat_session_id", new_column_name="chat_session_id"
-    )
-
-    op.drop_constraint("chat_session_pkey", "chat_session", type_="primary")
-    op.drop_column("chat_session", "id")
-    op.alter_column("chat_session", "new_id", new_column_name="id")
-
-    op.create_primary_key("chat_session_pkey", "chat_session", ["id"])
-
-    op.create_foreign_key(
-        "chat_message_chat_session_id_fkey",
-        "chat_message",
-        "chat_session",
-        ["chat_session_id"],
-        ["id"],
-        ondelete="CASCADE",
-    )
-
-
-def downgrade() -> None:
-    op.drop_constraint(
-        "chat_message_chat_session_id_fkey", "chat_message", type_="foreignkey"
-    )
-
-    op.add_column(
-        "chat_session",
-        sa.Column("old_id", sa.Integer, autoincrement=True, nullable=True),
-    )
-
-    op.execute("CREATE SEQUENCE chat_session_old_id_seq OWNED BY chat_session.old_id;")
-    op.execute(
-        "ALTER TABLE chat_session ALTER COLUMN old_id SET DEFAULT nextval('chat_session_old_id_seq');"
-    )
-
-    op.execute(
-        "UPDATE chat_session SET old_id = nextval('chat_session_old_id_seq') WHERE old_id IS NULL;"
-    )
-
-    op.alter_column("chat_session", "old_id", nullable=False)
-
-    op.drop_constraint("chat_session_pkey", "chat_session", type_="primary")
-    op.create_primary_key("chat_session_pkey", "chat_session", ["old_id"])
-
-    op.add_column(
-        "chat_message",
-        sa.Column("old_chat_session_id", sa.Integer, nullable=True),
-    )
-
-    op.execute(
-        """
-        UPDATE chat_message
-        SET old_chat_session_id = cs.old_id
-        FROM chat_session cs
-        WHERE chat_message.chat_session_id = cs.id;
-        """
-    )
-
-    op.drop_column("chat_message", "chat_session_id")
-    op.alter_column(
-        "chat_message", "old_chat_session_id", new_column_name="chat_session_id"
-    )
-
-    op.create_foreign_key(
-        "chat_message_chat_session_id_fkey",
-        "chat_message",
-        "chat_session",
-        ["chat_session_id"],
-        ["old_id"],
-        ondelete="CASCADE",
-    )
-
-    op.drop_column("chat_session", "id")
-    op.alter_column("chat_session", "old_id", new_column_name="id")
-
-    op.alter_column(
-        "chat_session",
-        "id",
-        type_=sa.Integer(),
-        existing_type=sa.Integer(),
-        existing_nullable=False,
-        existing_server_default=False,
-    )
-
-    # Rename the sequence
-    op.execute("ALTER SEQUENCE chat_session_old_id_seq RENAME TO chat_session_id_seq;")
-
-    # Update the default value to use the renamed sequence
-    op.alter_column(
-        "chat_session",
-        "id",
-        server_default=sa.text("nextval('chat_session_id_seq'::regclass)"),
-    )

backend/alembic/versions/b082fec533f0_make_last_attempt_status_nullable.py

@@ -31,12 +31,6 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
-    # First, update any null values to a default value
-    op.execute(
-        "UPDATE connector_credential_pair SET last_attempt_status = 'NOT_STARTED' WHERE last_attempt_status IS NULL"
-    )
-
-    # Then, make the column non-nullable
     op.alter_column(
         "connector_credential_pair",
         "last_attempt_status",

backend/alembic/versions/da4c21c69164_chosen_assistants_changed_to_jsonb.py

@@ -20,7 +20,7 @@ depends_on: None = None
 def upgrade() -> None:
     conn = op.get_bind()
     existing_ids_and_chosen_assistants = conn.execute(
-        sa.text('select id, chosen_assistants from "user"')
+        sa.text("select id, chosen_assistants from public.user")
     )
     op.drop_column(
         "user",
@@ -37,7 +37,7 @@ def upgrade() -> None:
     for id, chosen_assistants in existing_ids_and_chosen_assistants:
         conn.execute(
             sa.text(
-                'update "user" set chosen_assistants = :chosen_assistants where id = :id'
+                "update public.user set chosen_assistants = :chosen_assistants where id = :id"
             ),
             {"chosen_assistants": json.dumps(chosen_assistants), "id": id},
         )
@@ -46,7 +46,7 @@ def upgrade() -> None:
 def downgrade() -> None:
     conn = op.get_bind()
     existing_ids_and_chosen_assistants = conn.execute(
-        sa.text('select id, chosen_assistants from "user"')
+        sa.text("select id, chosen_assistants from public.user")
     )
     op.drop_column(
         "user",
@@ -59,7 +59,7 @@ def downgrade() -> None:
     for id, chosen_assistants in existing_ids_and_chosen_assistants:
         conn.execute(
             sa.text(
-                'update "user" set chosen_assistants = :chosen_assistants where id = :id'
+                "update public.user set chosen_assistants = :chosen_assistants where id = :id"
             ),
             {"chosen_assistants": chosen_assistants, "id": id},
         )

backend/alembic/versions/e4334d5b33ba_add_deployment_name_to_llmprovider.py

@@ -1,26 +0,0 @@
-"""add_deployment_name_to_llmprovider
-
-Revision ID: e4334d5b33ba
-Revises: ac5eaac849f9
-Create Date: 2024-10-04 09:52:34.896867
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "e4334d5b33ba"
-down_revision = "ac5eaac849f9"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "llm_provider", sa.Column("deployment_name", sa.String(), nullable=True)
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("llm_provider", "deployment_name")

backend/alembic_tenants/README.md

@@ -1,3 +0,0 @@
-These files are for public table migrations when operating with multi tenancy.
-
-If you are not a Danswer developer, you can ignore this directory entirely.

backend/alembic_tenants/env.py

@@ -1,111 +0,0 @@
-import asyncio
-from logging.config import fileConfig
-
-from sqlalchemy import pool
-from sqlalchemy.engine import Connection
-from sqlalchemy.ext.asyncio import create_async_engine
-from sqlalchemy.schema import SchemaItem
-
-from alembic import context
-from danswer.db.engine import build_connection_string
-from danswer.db.models import PublicBase
-
-# this is the Alembic Config object, which provides
-# access to the values within the .ini file in use.
-config = context.config
-
-# Interpret the config file for Python logging.
-# This line sets up loggers basically.
-if config.config_file_name is not None and config.attributes.get(
-    "configure_logger", True
-):
-    fileConfig(config.config_file_name)
-
-# add your model's MetaData object here
-# for 'autogenerate' support
-# from myapp import mymodel
-# target_metadata = mymodel.Base.metadata
-target_metadata = [PublicBase.metadata]
-
-# other values from the config, defined by the needs of env.py,
-# can be acquired:
-# my_important_option = config.get_main_option("my_important_option")
-# ... etc.
-
-EXCLUDE_TABLES = {"kombu_queue", "kombu_message"}
-
-
-def include_object(
-    object: SchemaItem,
-    name: str,
-    type_: str,
-    reflected: bool,
-    compare_to: SchemaItem | None,
-) -> bool:
-    if type_ == "table" and name in EXCLUDE_TABLES:
-        return False
-    return True
-
-
-def run_migrations_offline() -> None:
-    """Run migrations in 'offline' mode.
-
-    This configures the context with just a URL
-    and not an Engine, though an Engine is acceptable
-    here as well.  By skipping the Engine creation
-    we don't even need a DBAPI to be available.
-
-    Calls to context.execute() here emit the given string to the
-    script output.
-
-    """
-    url = build_connection_string()
-    context.configure(
-        url=url,
-        target_metadata=target_metadata,  # type: ignore
-        literal_binds=True,
-        dialect_opts={"paramstyle": "named"},
-    )
-
-    with context.begin_transaction():
-        context.run_migrations()
-
-
-def do_run_migrations(connection: Connection) -> None:
-    context.configure(
-        connection=connection,
-        target_metadata=target_metadata,  # type: ignore
-        include_object=include_object,
-    )  # type: ignore
-
-    with context.begin_transaction():
-        context.run_migrations()
-
-
-async def run_async_migrations() -> None:
-    """In this scenario we need to create an Engine
-    and associate a connection with the context.
-
-    """
-
-    connectable = create_async_engine(
-        build_connection_string(),
-        poolclass=pool.NullPool,
-    )
-
-    async with connectable.connect() as connection:
-        await connection.run_sync(do_run_migrations)
-
-    await connectable.dispose()
-
-
-def run_migrations_online() -> None:
-    """Run migrations in 'online' mode."""
-
-    asyncio.run(run_async_migrations())
-
-
-if context.is_offline_mode():
-    run_migrations_offline()
-else:
-    run_migrations_online()

backend/alembic_tenants/script.py.mako

@@ -1,24 +0,0 @@
-"""${message}
-
-Revision ID: ${up_revision}
-Revises: ${down_revision | comma,n}
-Create Date: ${create_date}
-
-"""
-from alembic import op
-import sqlalchemy as sa
-${imports if imports else ""}
-
-# revision identifiers, used by Alembic.
-revision = ${repr(up_revision)}
-down_revision = ${repr(down_revision)}
-branch_labels = ${repr(branch_labels)}
-depends_on = ${repr(depends_on)}
-
-
-def upgrade() -> None:
-    ${upgrades if upgrades else "pass"}
-
-
-def downgrade() -> None:
-    ${downgrades if downgrades else "pass"}

backend/alembic_tenants/versions/14a83a331951_create_usertenantmapping_table.py

@@ -1,24 +0,0 @@
-import sqlalchemy as sa
-
-from alembic import op
-
-# revision identifiers, used by Alembic.
-revision = "14a83a331951"
-down_revision = None
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.create_table(
-        "user_tenant_mapping",
-        sa.Column("email", sa.String(), nullable=False),
-        sa.Column("tenant_id", sa.String(), nullable=False),
-        sa.UniqueConstraint("email", "tenant_id", name="uq_user_tenant"),
-        sa.UniqueConstraint("email", name="uq_email"),
-        schema="public",
-    )
-
-
-def downgrade() -> None:
-    op.drop_table("user_tenant_mapping", schema="public")

backend/danswer/__init__.py

@@ -1,3 +1,3 @@
 import os
 
-__version__ = os.environ.get("DANSWER_VERSION", "") or "Development"
+__version__ = os.environ.get("DANSWER_VERSION", "") or "0.3-dev"

backend/danswer/access/models.py

@@ -70,12 +70,3 @@ class DocumentAccess(ExternalAccess):
             user_groups=set(user_groups),
             is_public=is_public,
         )
-
-
-default_public_access = DocumentAccess(
-    external_user_emails=set(),
-    external_user_group_ids=set(),
-    user_emails=set(),
-    user_groups=set(),
-    is_public=True,
-)

backend/danswer/auth/invited_users.py

@@ -9,7 +9,6 @@ from danswer.key_value_store.interface import KvKeyNotFoundError
 def get_invited_users() -> list[str]:
     try:
         store = get_kv_store()
-
         return cast(list, store.load(KV_USER_STORE_KEY))
     except KvKeyNotFoundError:
         return list()

backend/danswer/auth/schemas.py

@@ -34,7 +34,6 @@ class UserRead(schemas.BaseUser[uuid.UUID]):
 class UserCreate(schemas.BaseUserCreate):
     role: UserRole = UserRole.BASIC
     has_web_login: bool | None = True
-    tenant_id: str | None = None
 
 
 class UserUpdate(schemas.BaseUserUpdate):

backend/danswer/auth/users.py

@@ -5,23 +5,18 @@ from datetime import datetime
 from datetime import timezone
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
-from typing import Dict
-from typing import List
 from typing import Optional
 from typing import Tuple
 
 import jwt
 from email_validator import EmailNotValidError
-from email_validator import EmailUndeliverableError
 from email_validator import validate_email
 from fastapi import APIRouter
 from fastapi import Depends
 from fastapi import HTTPException
-from fastapi import Query
 from fastapi import Request
 from fastapi import Response
 from fastapi import status
-from fastapi.responses import RedirectResponse
 from fastapi.security import OAuth2PasswordRequestForm
 from fastapi_users import BaseUserManager
 from fastapi_users import exceptions
@@ -31,25 +26,11 @@ from fastapi_users import schemas
 from fastapi_users import UUIDIDMixin
 from fastapi_users.authentication import AuthenticationBackend
 from fastapi_users.authentication import CookieTransport
-from fastapi_users.authentication import JWTStrategy
 from fastapi_users.authentication import Strategy
 from fastapi_users.authentication.strategy.db import AccessTokenDatabase
 from fastapi_users.authentication.strategy.db import DatabaseStrategy
-from fastapi_users.exceptions import UserAlreadyExists
-from fastapi_users.jwt import decode_jwt
-from fastapi_users.jwt import generate_jwt
-from fastapi_users.jwt import SecretType
-from fastapi_users.manager import UserManagerDependency
 from fastapi_users.openapi import OpenAPIResponseType
-from fastapi_users.router.common import ErrorCode
-from fastapi_users.router.common import ErrorModel
 from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase
-from httpx_oauth.integrations.fastapi import OAuth2AuthorizeCallback
-from httpx_oauth.oauth2 import BaseOAuth2
-from httpx_oauth.oauth2 import OAuth2Token
-from pydantic import BaseModel
-from sqlalchemy import select
-from sqlalchemy.orm import attributes
 from sqlalchemy.orm import Session
 
 from danswer.auth.invited_users import get_invited_users
@@ -57,12 +38,11 @@ from danswer.auth.schemas import UserCreate
 from danswer.auth.schemas import UserRole
 from danswer.auth.schemas import UserUpdate
 from danswer.configs.app_configs import AUTH_TYPE
+from danswer.configs.app_configs import DATA_PLANE_SECRET
 from danswer.configs.app_configs import DISABLE_AUTH
-from danswer.configs.app_configs import DISABLE_VERIFICATION
 from danswer.configs.app_configs import EMAIL_FROM
-from danswer.configs.app_configs import MULTI_TENANT
+from danswer.configs.app_configs import EXPECTED_API_KEY
 from danswer.configs.app_configs import REQUIRE_EMAIL_VERIFICATION
-from danswer.configs.app_configs import SECRET_JWT_KEY
 from danswer.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS
 from danswer.configs.app_configs import SMTP_PASS
 from danswer.configs.app_configs import SMTP_PORT
@@ -80,23 +60,15 @@ from danswer.db.auth import get_access_token_db
 from danswer.db.auth import get_default_admin_user_emails
 from danswer.db.auth import get_user_count
 from danswer.db.auth import get_user_db
-from danswer.db.auth import SQLAlchemyUserAdminDB
-from danswer.db.engine import get_async_session_with_tenant
 from danswer.db.engine import get_session
-from danswer.db.engine import get_session_with_tenant
 from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.models import AccessToken
-from danswer.db.models import OAuthAccount
 from danswer.db.models import User
-from danswer.db.models import UserTenantMapping
 from danswer.db.users import get_user_by_email
 from danswer.utils.logger import setup_logger
 from danswer.utils.telemetry import optional_telemetry
 from danswer.utils.telemetry import RecordType
 from danswer.utils.variable_functionality import fetch_versioned_implementation
-from shared_configs.configs import CURRENT_TENANT_ID_CONTEXTVAR
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
-
 
 logger = setup_logger()
 
@@ -135,9 +107,7 @@ def get_display_email(email: str | None, space_less: bool = False) -> str:
 def user_needs_to_be_verified() -> bool:
     # all other auth types besides basic should require users to be
     # verified
-    return not DISABLE_VERIFICATION and (
-        AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
-    )
+    return AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
 
 
 def verify_email_is_invited(email: str) -> None:
@@ -148,10 +118,7 @@ def verify_email_is_invited(email: str) -> None:
     if not email:
         raise PermissionError("Email must be specified")
 
-    try:
-        email_info = validate_email(email)
-    except EmailUndeliverableError:
-        raise PermissionError("Email is not valid")
+    email_info = validate_email(email)  # can raise EmailNotValidError
 
     for email_whitelist in whitelist:
         try:
@@ -169,8 +136,8 @@ def verify_email_is_invited(email: str) -> None:
     raise PermissionError("User not on allowed user whitelist")
 
 
-def verify_email_in_whitelist(email: str, tenant_id: str | None = None) -> None:
-    with get_session_with_tenant(tenant_id) as db_session:
+def verify_email_in_whitelist(email: str) -> None:
+    with Session(get_sqlalchemy_engine()) as db_session:
         if not get_user_by_email(email, db_session):
             verify_email_is_invited(email)
 
@@ -190,20 +157,6 @@ def verify_email_domain(email: str) -> None:
             )
 
 
-def get_tenant_id_for_email(email: str) -> str:
-    if not MULTI_TENANT:
-        return POSTGRES_DEFAULT_SCHEMA
-    # Implement logic to get tenant_id from the mapping table
-    with Session(get_sqlalchemy_engine()) as db_session:
-        result = db_session.execute(
-            select(UserTenantMapping.tenant_id).where(UserTenantMapping.email == email)
-        )
-        tenant_id = result.scalar_one_or_none()
-    if tenant_id is None:
-        raise exceptions.UserNotExists()
-    return tenant_id
-
-
 def send_user_verification_email(
     user_email: str,
     token: str,
@@ -238,85 +191,35 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         safe: bool = False,
         request: Optional[Request] = None,
     ) -> User:
+        verify_email_is_invited(user_create.email)
+        verify_email_domain(user_create.email)
+        if hasattr(user_create, "role"):
+            user_count = await get_user_count()
+            if user_count == 0 or user_create.email in get_default_admin_user_emails():
+                user_create.role = UserRole.ADMIN
+            else:
+                user_create.role = UserRole.BASIC
+        user = None
         try:
-            tenant_id = (
-                get_tenant_id_for_email(user_create.email)
-                if MULTI_TENANT
-                else POSTGRES_DEFAULT_SCHEMA
-            )
-        except exceptions.UserNotExists:
-            raise HTTPException(status_code=401, detail="User not found")
-
-        if not tenant_id:
-            raise HTTPException(
-                status_code=401, detail="User does not belong to an organization"
-            )
-
-        async with get_async_session_with_tenant(tenant_id) as db_session:
-            token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
-
-            verify_email_is_invited(user_create.email)
-            verify_email_domain(user_create.email)
-            if MULTI_TENANT:
-                tenant_user_db = SQLAlchemyUserAdminDB(db_session, User, OAuthAccount)
-                self.user_db = tenant_user_db
-                self.database = tenant_user_db
-
-            if hasattr(user_create, "role"):
-                user_count = await get_user_count()
-                if (
-                    user_count == 0
-                    or user_create.email in get_default_admin_user_emails()
-                ):
-                    user_create.role = UserRole.ADMIN
-                else:
-                    user_create.role = UserRole.BASIC
-            user = None
-            try:
-                user = await super().create(user_create, safe=safe, request=request)  # type: ignore
-            except exceptions.UserAlreadyExists:
-                user = await self.get_by_email(user_create.email)
-                # Handle case where user has used product outside of web and is now creating an account through web
-                if (
-                    not user.has_web_login
-                    and hasattr(user_create, "has_web_login")
-                    and user_create.has_web_login
-                ):
-                    user_update = UserUpdate(
-                        password=user_create.password,
-                        has_web_login=True,
-                        role=user_create.role,
-                        is_verified=user_create.is_verified,
-                    )
-                    user = await self.update(user_update, user)
-                else:
-                    raise exceptions.UserAlreadyExists()
-
-            CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
-            return user
-
-    async def on_after_login(
-        self,
-        user: User,
-        request: Request | None = None,
-        response: Response | None = None,
-    ) -> None:
-        if response is None or not MULTI_TENANT:
-            return
-
-        tenant_id = get_tenant_id_for_email(user.email)
-
-        tenant_token = jwt.encode(
-            {"tenant_id": tenant_id}, SECRET_JWT_KEY, algorithm="HS256"
-        )
-
-        response.set_cookie(
-            key="tenant_details",
-            value=tenant_token,
-            httponly=True,
-            secure=WEB_DOMAIN.startswith("https"),
-            samesite="lax",
-        )
+            user = await super().create(user_create, safe=safe, request=request)  # type: ignore
+        except exceptions.UserAlreadyExists:
+            user = await self.get_by_email(user_create.email)
+            # Handle case where user has used product outside of web and is now creating an account through web
+            if (
+                not user.has_web_login
+                and hasattr(user_create, "has_web_login")
+                and user_create.has_web_login
+            ):
+                user_update = UserUpdate(
+                    password=user_create.password,
+                    has_web_login=True,
+                    role=user_create.role,
+                    is_verified=user_create.is_verified,
+                )
+                user = await self.update(user_update, user)
+            else:
+                raise exceptions.UserAlreadyExists()
+        return user
 
     async def oauth_callback(
         self: "BaseUserManager[models.UOAP, models.ID]",
@@ -331,114 +234,45 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         associate_by_email: bool = False,
         is_verified_by_default: bool = False,
     ) -> models.UOAP:
-        # Get tenant_id from mapping table
-        try:
-            tenant_id = (
-                get_tenant_id_for_email(account_email)
-                if MULTI_TENANT
-                else POSTGRES_DEFAULT_SCHEMA
+        verify_email_in_whitelist(account_email)
+        verify_email_domain(account_email)
+
+        user = await super().oauth_callback(  # type: ignore
+            oauth_name=oauth_name,
+            access_token=access_token,
+            account_id=account_id,
+            account_email=account_email,
+            expires_at=expires_at,
+            refresh_token=refresh_token,
+            request=request,
+            associate_by_email=associate_by_email,
+            is_verified_by_default=is_verified_by_default,
+        )
+
+        # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to
+        # re-authenticate that frequently, so by default this is disabled
+        if expires_at and TRACK_EXTERNAL_IDP_EXPIRY:
+            oidc_expiry = datetime.fromtimestamp(expires_at, tz=timezone.utc)
+            await self.user_db.update(user, update_dict={"oidc_expiry": oidc_expiry})
+
+        # this is needed if an organization goes from `TRACK_EXTERNAL_IDP_EXPIRY=true` to `false`
+        # otherwise, the oidc expiry will always be old, and the user will never be able to login
+        if user.oidc_expiry and not TRACK_EXTERNAL_IDP_EXPIRY:
+            await self.user_db.update(user, update_dict={"oidc_expiry": None})
+
+        # Handle case where user has used product outside of web and is now creating an account through web
+        if not user.has_web_login:
+            await self.user_db.update(
+                user,
+                update_dict={
+                    "is_verified": is_verified_by_default,
+                    "has_web_login": True,
+                },
             )
-        except exceptions.UserNotExists:
-            raise HTTPException(status_code=401, detail="User not found")
+            user.is_verified = is_verified_by_default
+            user.has_web_login = True
 
-        if not tenant_id:
-            raise HTTPException(status_code=401, detail="User not found")
-
-        token = None
-        async with get_async_session_with_tenant(tenant_id) as db_session:
-            token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
-
-            verify_email_in_whitelist(account_email, tenant_id)
-            verify_email_domain(account_email)
-
-            if MULTI_TENANT:
-                tenant_user_db = SQLAlchemyUserAdminDB(db_session, User, OAuthAccount)
-                self.user_db = tenant_user_db
-                self.database = tenant_user_db  # type: ignore
-
-            oauth_account_dict = {
-                "oauth_name": oauth_name,
-                "access_token": access_token,
-                "account_id": account_id,
-                "account_email": account_email,
-                "expires_at": expires_at,
-                "refresh_token": refresh_token,
-            }
-
-            try:
-                # Attempt to get user by OAuth account
-                user = await self.get_by_oauth_account(oauth_name, account_id)
-
-            except exceptions.UserNotExists:
-                try:
-                    # Attempt to get user by email
-                    user = await self.get_by_email(account_email)
-                    if not associate_by_email:
-                        raise exceptions.UserAlreadyExists()
-
-                    user = await self.user_db.add_oauth_account(
-                        user, oauth_account_dict
-                    )
-
-                    # If user not found by OAuth account or email, create a new user
-                except exceptions.UserNotExists:
-                    password = self.password_helper.generate()
-                    user_dict = {
-                        "email": account_email,
-                        "hashed_password": self.password_helper.hash(password),
-                        "is_verified": is_verified_by_default,
-                    }
-
-                    user = await self.user_db.create(user_dict)
-                    user = await self.user_db.add_oauth_account(
-                        user, oauth_account_dict
-                    )
-                    await self.on_after_register(user, request)
-
-            else:
-                for existing_oauth_account in user.oauth_accounts:
-                    if (
-                        existing_oauth_account.account_id == account_id
-                        and existing_oauth_account.oauth_name == oauth_name
-                    ):
-                        user = await self.user_db.update_oauth_account(
-                            user, existing_oauth_account, oauth_account_dict
-                        )
-
-            # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to
-            # re-authenticate that frequently, so by default this is disabled
-
-            if expires_at and TRACK_EXTERNAL_IDP_EXPIRY:
-                oidc_expiry = datetime.fromtimestamp(expires_at, tz=timezone.utc)
-                await self.user_db.update(
-                    user, update_dict={"oidc_expiry": oidc_expiry}
-                )
-
-            # Handle case where user has used product outside of web and is now creating an account through web
-            if not user.has_web_login:  # type: ignore
-                await self.user_db.update(
-                    user,
-                    {
-                        "is_verified": is_verified_by_default,
-                        "has_web_login": True,
-                    },
-                )
-                user.is_verified = is_verified_by_default
-                user.has_web_login = True  # type: ignore
-
-            # this is needed if an organization goes from `TRACK_EXTERNAL_IDP_EXPIRY=true` to `false`
-            # otherwise, the oidc expiry will always be old, and the user will never be able to login
-            if (
-                user.oidc_expiry is not None  # type: ignore
-                and not TRACK_EXTERNAL_IDP_EXPIRY
-            ):
-                await self.user_db.update(user, {"oidc_expiry": None})
-                user.oidc_expiry = None  # type: ignore
-
-            if token:
-                CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
-
-            return user
+        return user
 
     async def on_after_register(
         self, user: User, request: Optional[Request] = None
@@ -469,50 +303,28 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
     async def authenticate(
         self, credentials: OAuth2PasswordRequestForm
     ) -> Optional[User]:
-        email = credentials.username
-
-        # Get tenant_id from mapping table
-        tenant_id = get_tenant_id_for_email(email)
-        if not tenant_id:
-            # User not found in mapping
+        try:
+            user = await self.get_by_email(credentials.username)
+        except exceptions.UserNotExists:
             self.password_helper.hash(credentials.password)
             return None
 
-        # Create a tenant-specific session
-        async with get_async_session_with_tenant(tenant_id) as tenant_session:
-            tenant_user_db: SQLAlchemyUserDatabase = SQLAlchemyUserDatabase(
-                tenant_session, User
+        if not user.has_web_login:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="NO_WEB_LOGIN_AND_HAS_NO_PASSWORD",
             )
-            self.user_db = tenant_user_db
 
-            # Proceed with authentication
-            try:
-                user = await self.get_by_email(email)
+        verified, updated_password_hash = self.password_helper.verify_and_update(
+            credentials.password, user.hashed_password
+        )
+        if not verified:
+            return None
 
-            except exceptions.UserNotExists:
-                self.password_helper.hash(credentials.password)
-                return None
+        if updated_password_hash is not None:
+            await self.user_db.update(user, {"hashed_password": updated_password_hash})
 
-            has_web_login = attributes.get_attribute(user, "has_web_login")
-
-            if not has_web_login:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail="NO_WEB_LOGIN_AND_HAS_NO_PASSWORD",
-                )
-
-            verified, updated_password_hash = self.password_helper.verify_and_update(
-                credentials.password, user.hashed_password
-            )
-            if not verified:
-                return None
-
-            if updated_password_hash is not None:
-                await self.user_db.update(
-                    user, {"hashed_password": updated_password_hash}
-                )
-
-            return user
+        return user
 
 
 async def get_user_manager(
@@ -527,26 +339,20 @@ cookie_transport = CookieTransport(
 )
 
 
-def get_jwt_strategy() -> JWTStrategy:
-    return JWTStrategy(
-        secret=USER_AUTH_SECRET,
-        lifetime_seconds=SESSION_EXPIRE_TIME_SECONDS,
-    )
-
-
 def get_database_strategy(
     access_token_db: AccessTokenDatabase[AccessToken] = Depends(get_access_token_db),
 ) -> DatabaseStrategy:
-    return DatabaseStrategy(
+    strategy = DatabaseStrategy(
         access_token_db, lifetime_seconds=SESSION_EXPIRE_TIME_SECONDS  # type: ignore
     )
+    return strategy
 
 
 auth_backend = AuthenticationBackend(
-    name="jwt" if MULTI_TENANT else "database",
+    name="database",
     transport=cookie_transport,
-    get_strategy=get_jwt_strategy if MULTI_TENANT else get_database_strategy,  # type: ignore
-)  # type: ignore
+    get_strategy=get_database_strategy,
+)
 
 
 class FastAPIUserWithLogoutRouter(FastAPIUsers[models.UP, models.ID]):
@@ -560,11 +366,9 @@ class FastAPIUserWithLogoutRouter(FastAPIUsers[models.UP, models.ID]):
         This way the login router does not need to be included
         """
         router = APIRouter()
-
         get_current_user_token = self.authenticator.current_user_token(
             active=True, verified=requires_verification
         )
-
         logout_responses: OpenAPIResponseType = {
             **{
                 status.HTTP_401_UNAUTHORIZED: {
@@ -611,8 +415,8 @@ async def optional_user_(
 
 async def optional_user(
     request: Request,
-    db_session: Session = Depends(get_session),
     user: User | None = Depends(optional_fastapi_current_user),
+    db_session: Session = Depends(get_session),
 ) -> User | None:
     versioned_fetch_user = fetch_versioned_implementation(
         "danswer.auth.users", "optional_user_"
@@ -705,184 +509,26 @@ def get_default_admin_user_emails_() -> list[str]:
     return []
 
 
-STATE_TOKEN_AUDIENCE = "fastapi-users:oauth-state"
+async def control_plane_dep(request: Request) -> None:
+    api_key = request.headers.get("X-API-KEY")
+    if api_key != EXPECTED_API_KEY:
+        logger.warning("Invalid API key")
+        raise HTTPException(status_code=401, detail="Invalid API key")
 
+    auth_header = request.headers.get("Authorization")
+    if not auth_header or not auth_header.startswith("Bearer "):
+        logger.warning("Invalid authorization header")
+        raise HTTPException(status_code=401, detail="Invalid authorization header")
 
-class OAuth2AuthorizeResponse(BaseModel):
-    authorization_url: str
-
-
-def generate_state_token(
-    data: Dict[str, str], secret: SecretType, lifetime_seconds: int = 3600
-) -> str:
-    data["aud"] = STATE_TOKEN_AUDIENCE
-
-    return generate_jwt(data, secret, lifetime_seconds)
-
-
-# refer to https://github.com/fastapi-users/fastapi-users/blob/42ddc241b965475390e2bce887b084152ae1a2cd/fastapi_users/fastapi_users.py#L91
-
-
-def create_danswer_oauth_router(
-    oauth_client: BaseOAuth2,
-    backend: AuthenticationBackend,
-    state_secret: SecretType,
-    redirect_url: Optional[str] = None,
-    associate_by_email: bool = False,
-    is_verified_by_default: bool = False,
-) -> APIRouter:
-    return get_oauth_router(
-        oauth_client,
-        backend,
-        get_user_manager,
-        state_secret,
-        redirect_url,
-        associate_by_email,
-        is_verified_by_default,
-    )
-
-
-def get_oauth_router(
-    oauth_client: BaseOAuth2,
-    backend: AuthenticationBackend,
-    get_user_manager: UserManagerDependency[models.UP, models.ID],
-    state_secret: SecretType,
-    redirect_url: Optional[str] = None,
-    associate_by_email: bool = False,
-    is_verified_by_default: bool = False,
-) -> APIRouter:
-    """Generate a router with the OAuth routes."""
-    router = APIRouter()
-    callback_route_name = f"oauth:{oauth_client.name}.{backend.name}.callback"
-
-    if redirect_url is not None:
-        oauth2_authorize_callback = OAuth2AuthorizeCallback(
-            oauth_client,
-            redirect_url=redirect_url,
-        )
-    else:
-        oauth2_authorize_callback = OAuth2AuthorizeCallback(
-            oauth_client,
-            route_name=callback_route_name,
-        )
-
-    @router.get(
-        "/authorize",
-        name=f"oauth:{oauth_client.name}.{backend.name}.authorize",
-        response_model=OAuth2AuthorizeResponse,
-    )
-    async def authorize(
-        request: Request, scopes: List[str] = Query(None)
-    ) -> OAuth2AuthorizeResponse:
-        if redirect_url is not None:
-            authorize_redirect_url = redirect_url
-        else:
-            authorize_redirect_url = str(request.url_for(callback_route_name))
-
-        next_url = request.query_params.get("next", "/")
-        state_data: Dict[str, str] = {"next_url": next_url}
-        state = generate_state_token(state_data, state_secret)
-        authorization_url = await oauth_client.get_authorization_url(
-            authorize_redirect_url,
-            state,
-            scopes,
-        )
-
-        return OAuth2AuthorizeResponse(authorization_url=authorization_url)
-
-    @router.get(
-        "/callback",
-        name=callback_route_name,
-        description="The response varies based on the authentication backend used.",
-        responses={
-            status.HTTP_400_BAD_REQUEST: {
-                "model": ErrorModel,
-                "content": {
-                    "application/json": {
-                        "examples": {
-                            "INVALID_STATE_TOKEN": {
-                                "summary": "Invalid state token.",
-                                "value": None,
-                            },
-                            ErrorCode.LOGIN_BAD_CREDENTIALS: {
-                                "summary": "User is inactive.",
-                                "value": {"detail": ErrorCode.LOGIN_BAD_CREDENTIALS},
-                            },
-                        }
-                    }
-                },
-            },
-        },
-    )
-    async def callback(
-        request: Request,
-        access_token_state: Tuple[OAuth2Token, str] = Depends(
-            oauth2_authorize_callback
-        ),
-        user_manager: BaseUserManager[models.UP, models.ID] = Depends(get_user_manager),
-        strategy: Strategy[models.UP, models.ID] = Depends(backend.get_strategy),
-    ) -> RedirectResponse:
-        token, state = access_token_state
-        account_id, account_email = await oauth_client.get_id_email(
-            token["access_token"]
-        )
-
-        if account_email is None:
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST,
-                detail=ErrorCode.OAUTH_NOT_AVAILABLE_EMAIL,
-            )
-
-        try:
-            state_data = decode_jwt(state, state_secret, [STATE_TOKEN_AUDIENCE])
-        except jwt.DecodeError:
-            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
-
-        next_url = state_data.get("next_url", "/")
-
-        # Authenticate user
-        try:
-            user = await user_manager.oauth_callback(
-                oauth_client.name,
-                token["access_token"],
-                account_id,
-                account_email,
-                token.get("expires_at"),
-                token.get("refresh_token"),
-                request,
-                associate_by_email=associate_by_email,
-                is_verified_by_default=is_verified_by_default,
-            )
-        except UserAlreadyExists:
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST,
-                detail=ErrorCode.OAUTH_USER_ALREADY_EXISTS,
-            )
-
-        if not user.is_active:
-            raise HTTPException(
-                status_code=status.HTTP_400_BAD_REQUEST,
-                detail=ErrorCode.LOGIN_BAD_CREDENTIALS,
-            )
-
-        # Login user
-        response = await backend.login(strategy, user)
-        await user_manager.on_after_login(user, request, response)
-
-        # Prepare redirect response
-        redirect_response = RedirectResponse(next_url, status_code=302)
-
-        # Copy headers and other attributes from 'response' to 'redirect_response'
-        for header_name, header_value in response.headers.items():
-            redirect_response.headers[header_name] = header_value
-
-        if hasattr(response, "body"):
-            redirect_response.body = response.body
-        if hasattr(response, "status_code"):
-            redirect_response.status_code = response.status_code
-        if hasattr(response, "media_type"):
-            redirect_response.media_type = response.media_type
-
-        return redirect_response
-
-    return router
+    token = auth_header.split(" ")[1]
+    try:
+        payload = jwt.decode(token, DATA_PLANE_SECRET, algorithms=["HS256"])
+        if payload.get("scope") != "tenant:create":
+            logger.warning("Insufficient permissions")
+            raise HTTPException(status_code=403, detail="Insufficient permissions")
+    except jwt.ExpiredSignatureError:
+        logger.warning("Token has expired")
+        raise HTTPException(status_code=401, detail="Token has expired")
+    except jwt.InvalidTokenError:
+        logger.warning("Invalid token")
+        raise HTTPException(status_code=401, detail="Invalid token")

backend/danswer/background/celery/apps/app_base.py

@@ -1,310 +0,0 @@
-import logging
-import multiprocessing
-import time
-from typing import Any
-
-import sentry_sdk
-from celery import Task
-from celery.app import trace
-from celery.exceptions import WorkerShutdown
-from celery.states import READY_STATES
-from celery.utils.log import get_task_logger
-from celery.worker import strategy  # type: ignore
-from sentry_sdk.integrations.celery import CeleryIntegration
-
-from danswer.background.celery.apps.task_formatters import CeleryTaskColoredFormatter
-from danswer.background.celery.apps.task_formatters import CeleryTaskPlainFormatter
-from danswer.background.celery.celery_redis import RedisConnectorCredentialPair
-from danswer.background.celery.celery_redis import RedisConnectorDeletion
-from danswer.background.celery.celery_redis import RedisConnectorPruning
-from danswer.background.celery.celery_redis import RedisDocumentSet
-from danswer.background.celery.celery_redis import RedisUserGroup
-from danswer.background.celery.celery_utils import celery_is_worker_primary
-from danswer.configs.constants import DanswerRedisLocks
-from danswer.db.engine import get_all_tenant_ids
-from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import ColoredFormatter
-from danswer.utils.logger import PlainFormatter
-from danswer.utils.logger import setup_logger
-from shared_configs.configs import SENTRY_DSN
-
-
-logger = setup_logger()
-
-task_logger = get_task_logger(__name__)
-
-if SENTRY_DSN:
-    sentry_sdk.init(
-        dsn=SENTRY_DSN,
-        integrations=[CeleryIntegration()],
-        traces_sample_rate=0.1,
-    )
-    logger.info("Sentry initialized")
-else:
-    logger.debug("Sentry DSN not provided, skipping Sentry initialization")
-
-
-def on_task_prerun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    **kwds: Any,
-) -> None:
-    pass
-
-
-def on_task_postrun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict[str, Any] | None = None,
-    retval: Any | None = None,
-    state: str | None = None,
-    **kwds: Any,
-) -> None:
-    """We handle this signal in order to remove completed tasks
-    from their respective tasksets. This allows us to track the progress of document set
-    and user group syncs.
-
-    This function runs after any task completes (both success and failure)
-    Note that this signal does not fire on a task that failed to complete and is going
-    to be retried.
-
-    This also does not fire if a worker with acks_late=False crashes (which all of our
-    long running workers are)
-    """
-    if not task:
-        return
-
-    task_logger.debug(f"Task {task.name} (ID: {task_id}) completed with state: {state}")
-
-    if state not in READY_STATES:
-        return
-
-    if not task_id:
-        return
-
-    # Get tenant_id directly from kwargs- each celery task has a tenant_id kwarg
-    if not kwargs:
-        logger.error(f"Task {task.name} (ID: {task_id}) is missing kwargs")
-        tenant_id = None
-    else:
-        tenant_id = kwargs.get("tenant_id")
-
-    task_logger.debug(
-        f"Task {task.name} (ID: {task_id}) completed with state: {state} "
-        f"{f'for tenant_id={tenant_id}' if tenant_id else ''}"
-    )
-
-    r = get_redis_client(tenant_id=tenant_id)
-
-    if task_id.startswith(RedisConnectorCredentialPair.PREFIX):
-        r.srem(RedisConnectorCredentialPair.get_taskset_key(), task_id)
-        return
-
-    if task_id.startswith(RedisDocumentSet.PREFIX):
-        document_set_id = RedisDocumentSet.get_id_from_task_id(task_id)
-        if document_set_id is not None:
-            rds = RedisDocumentSet(int(document_set_id))
-            r.srem(rds.taskset_key, task_id)
-        return
-
-    if task_id.startswith(RedisUserGroup.PREFIX):
-        usergroup_id = RedisUserGroup.get_id_from_task_id(task_id)
-        if usergroup_id is not None:
-            rug = RedisUserGroup(int(usergroup_id))
-            r.srem(rug.taskset_key, task_id)
-        return
-
-    if task_id.startswith(RedisConnectorDeletion.PREFIX):
-        cc_pair_id = RedisConnectorDeletion.get_id_from_task_id(task_id)
-        if cc_pair_id is not None:
-            rcd = RedisConnectorDeletion(int(cc_pair_id))
-            r.srem(rcd.taskset_key, task_id)
-        return
-
-    if task_id.startswith(RedisConnectorPruning.SUBTASK_PREFIX):
-        cc_pair_id = RedisConnectorPruning.get_id_from_task_id(task_id)
-        if cc_pair_id is not None:
-            rcp = RedisConnectorPruning(int(cc_pair_id))
-            r.srem(rcp.taskset_key, task_id)
-        return
-
-
-def on_celeryd_init(sender: Any = None, conf: Any = None, **kwargs: Any) -> None:
-    """The first signal sent on celery worker startup"""
-    multiprocessing.set_start_method("spawn")  # fork is unsafe, set to spawn
-
-
-def wait_for_redis(sender: Any, **kwargs: Any) -> None:
-    r = get_redis_client(tenant_id=None)
-
-    WAIT_INTERVAL = 5
-    WAIT_LIMIT = 60
-
-    time_start = time.monotonic()
-    logger.info("Redis: Readiness check starting.")
-    while True:
-        try:
-            if r.ping():
-                break
-        except Exception:
-            pass
-
-        time_elapsed = time.monotonic() - time_start
-        logger.info(
-            f"Redis: Ping failed. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
-        )
-        if time_elapsed > WAIT_LIMIT:
-            msg = (
-                f"Redis: Readiness check did not succeed within the timeout "
-                f"({WAIT_LIMIT} seconds). Exiting..."
-            )
-            logger.error(msg)
-            raise WorkerShutdown(msg)
-
-        time.sleep(WAIT_INTERVAL)
-
-    logger.info("Redis: Readiness check succeeded. Continuing...")
-    return
-
-
-def on_secondary_worker_init(sender: Any, **kwargs: Any) -> None:
-    WAIT_INTERVAL = 5
-    WAIT_LIMIT = 60
-
-    logger.info("Running as a secondary celery worker.")
-    logger.info("Waiting for all tenant primary workers to be ready...")
-    time_start = time.monotonic()
-
-    while True:
-        tenant_ids = get_all_tenant_ids()
-        # Check if we have a primary worker lock for each tenant
-        all_tenants_ready = all(
-            get_redis_client(tenant_id=tenant_id).exists(
-                DanswerRedisLocks.PRIMARY_WORKER
-            )
-            for tenant_id in tenant_ids
-        )
-
-        if all_tenants_ready:
-            break
-
-        time_elapsed = time.monotonic() - time_start
-        ready_tenants = sum(
-            1
-            for tenant_id in tenant_ids
-            if get_redis_client(tenant_id=tenant_id).exists(
-                DanswerRedisLocks.PRIMARY_WORKER
-            )
-        )
-
-        logger.info(
-            f"Not all tenant primary workers are ready yet. "
-            f"Ready tenants: {ready_tenants}/{len(tenant_ids)} "
-            f"elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
-        )
-
-        if time_elapsed > WAIT_LIMIT:
-            msg = (
-                f"Not all tenant primary workers were ready within the timeout "
-                f"({WAIT_LIMIT} seconds). Exiting..."
-            )
-            logger.error(msg)
-            raise WorkerShutdown(msg)
-
-        time.sleep(WAIT_INTERVAL)
-
-    logger.info("All tenant primary workers are ready. Continuing...")
-    return
-
-
-def on_worker_ready(sender: Any, **kwargs: Any) -> None:
-    task_logger.info("worker_ready signal received.")
-
-
-def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
-    if not celery_is_worker_primary(sender):
-        return
-
-    if not hasattr(sender, "primary_worker_locks"):
-        return
-
-    for tenant_id, lock in sender.primary_worker_locks.items():
-        try:
-            if lock and lock.owned():
-                logger.debug(f"Attempting to release lock for tenant {tenant_id}")
-                try:
-                    lock.release()
-                    logger.debug(f"Successfully released lock for tenant {tenant_id}")
-                except Exception as e:
-                    logger.error(
-                        f"Failed to release lock for tenant {tenant_id}. Error: {str(e)}"
-                    )
-                finally:
-                    sender.primary_worker_locks[tenant_id] = None
-        except Exception as e:
-            logger.error(
-                f"Error checking lock status for tenant {tenant_id}. Error: {str(e)}"
-            )
-
-
-def on_setup_logging(
-    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
-) -> None:
-    # TODO: could unhardcode format and colorize and accept these as options from
-    # celery's config
-
-    # reformats the root logger
-    root_logger = logging.getLogger()
-
-    root_handler = logging.StreamHandler()  # Set up a handler for the root logger
-    root_formatter = ColoredFormatter(
-        "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
-        datefmt="%m/%d/%Y %I:%M:%S %p",
-    )
-    root_handler.setFormatter(root_formatter)
-    root_logger.addHandler(root_handler)  # Apply the handler to the root logger
-
-    if logfile:
-        root_file_handler = logging.FileHandler(logfile)
-        root_file_formatter = PlainFormatter(
-            "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
-            datefmt="%m/%d/%Y %I:%M:%S %p",
-        )
-        root_file_handler.setFormatter(root_file_formatter)
-        root_logger.addHandler(root_file_handler)
-
-    root_logger.setLevel(loglevel)
-
-    # reformats celery's task logger
-    task_formatter = CeleryTaskColoredFormatter(
-        "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
-        datefmt="%m/%d/%Y %I:%M:%S %p",
-    )
-    task_handler = logging.StreamHandler()  # Set up a handler for the task logger
-    task_handler.setFormatter(task_formatter)
-    task_logger.addHandler(task_handler)  # Apply the handler to the task logger
-
-    if logfile:
-        task_file_handler = logging.FileHandler(logfile)
-        task_file_formatter = CeleryTaskPlainFormatter(
-            "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
-            datefmt="%m/%d/%Y %I:%M:%S %p",
-        )
-        task_file_handler.setFormatter(task_file_formatter)
-        task_logger.addHandler(task_file_handler)
-
-    task_logger.setLevel(loglevel)
-    task_logger.propagate = False
-
-    # hide celery task received spam
-    # e.g. "Task check_for_pruning[a1e96171-0ba8-4e00-887b-9fbf7442eab3] received"
-    strategy.logger.setLevel(logging.WARNING)
-
-    # hide celery task succeeded/failed spam
-    # e.g. "Task check_for_pruning[a1e96171-0ba8-4e00-887b-9fbf7442eab3] succeeded in 0.03137450001668185s: None"
-    trace.logger.setLevel(logging.WARNING)

backend/danswer/background/celery/apps/beat.py

@@ -1,99 +0,0 @@
-from datetime import timedelta
-from typing import Any
-
-from celery import Celery
-from celery import signals
-from celery.signals import beat_init
-
-import danswer.background.celery.apps.app_base as app_base
-from danswer.configs.constants import DanswerCeleryPriority
-from danswer.configs.constants import POSTGRES_CELERY_BEAT_APP_NAME
-from danswer.db.engine import get_all_tenant_ids
-from danswer.db.engine import SqlEngine
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
-
-celery_app = Celery(__name__)
-celery_app.config_from_object("danswer.background.celery.configs.beat")
-
-
-@beat_init.connect
-def on_beat_init(sender: Any, **kwargs: Any) -> None:
-    logger.info("beat_init signal received.")
-    SqlEngine.set_app_name(POSTGRES_CELERY_BEAT_APP_NAME)
-    SqlEngine.init_engine(pool_size=2, max_overflow=0)
-    app_base.wait_for_redis(sender, **kwargs)
-
-
-@signals.setup_logging.connect
-def on_setup_logging(
-    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
-) -> None:
-    app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
-
-
-#####
-# Celery Beat (Periodic Tasks) Settings
-#####
-
-tenant_ids = get_all_tenant_ids()
-
-tasks_to_schedule = [
-    {
-        "name": "check-for-vespa-sync",
-        "task": "check_for_vespa_sync_task",
-        "schedule": timedelta(seconds=5),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-connector-deletion",
-        "task": "check_for_connector_deletion_task",
-        "schedule": timedelta(seconds=60),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-indexing",
-        "task": "check_for_indexing",
-        "schedule": timedelta(seconds=10),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-prune",
-        "task": "check_for_pruning",
-        "schedule": timedelta(seconds=10),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "kombu-message-cleanup",
-        "task": "kombu_message_cleanup_task",
-        "schedule": timedelta(seconds=3600),
-        "options": {"priority": DanswerCeleryPriority.LOWEST},
-    },
-    {
-        "name": "monitor-vespa-sync",
-        "task": "monitor_vespa_sync",
-        "schedule": timedelta(seconds=5),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-]
-
-# Build the celery beat schedule dynamically
-beat_schedule = {}
-
-for tenant_id in tenant_ids:
-    for task in tasks_to_schedule:
-        task_name = f"{task['name']}-{tenant_id}"  # Unique name for each scheduled task
-        beat_schedule[task_name] = {
-            "task": task["task"],
-            "schedule": task["schedule"],
-            "options": task["options"],
-            "kwargs": {"tenant_id": tenant_id},  # Must pass tenant_id as an argument
-        }
-
-# Include any existing beat schedules
-existing_beat_schedule = celery_app.conf.beat_schedule or {}
-beat_schedule.update(existing_beat_schedule)
-
-# Update the Celery app configuration once
-celery_app.conf.beat_schedule = beat_schedule

backend/danswer/background/celery/apps/heavy.py

@@ -1,88 +0,0 @@
-import multiprocessing
-from typing import Any
-
-from celery import Celery
-from celery import signals
-from celery import Task
-from celery.signals import celeryd_init
-from celery.signals import worker_init
-from celery.signals import worker_ready
-from celery.signals import worker_shutdown
-
-import danswer.background.celery.apps.app_base as app_base
-from danswer.configs.constants import POSTGRES_CELERY_WORKER_HEAVY_APP_NAME
-from danswer.db.engine import SqlEngine
-from danswer.utils.logger import setup_logger
-
-
-logger = setup_logger()
-
-celery_app = Celery(__name__)
-celery_app.config_from_object("danswer.background.celery.configs.heavy")
-
-
-@signals.task_prerun.connect
-def on_task_prerun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)
-
-
-@signals.task_postrun.connect
-def on_task_postrun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    retval: Any | None = None,
-    state: str | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)
-
-
-@celeryd_init.connect
-def on_celeryd_init(sender: Any = None, conf: Any = None, **kwargs: Any) -> None:
-    app_base.on_celeryd_init(sender, conf, **kwargs)
-
-
-@worker_init.connect
-def on_worker_init(sender: Any, **kwargs: Any) -> None:
-    logger.info("worker_init signal received.")
-    logger.info(f"Multiprocessing start method: {multiprocessing.get_start_method()}")
-
-    SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_HEAVY_APP_NAME)
-    SqlEngine.init_engine(pool_size=8, max_overflow=0)
-
-    app_base.wait_for_redis(sender, **kwargs)
-    app_base.on_secondary_worker_init(sender, **kwargs)
-
-
-@worker_ready.connect
-def on_worker_ready(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_ready(sender, **kwargs)
-
-
-@worker_shutdown.connect
-def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_shutdown(sender, **kwargs)
-
-
-@signals.setup_logging.connect
-def on_setup_logging(
-    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
-) -> None:
-    app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
-
-
-celery_app.autodiscover_tasks(
-    [
-        "danswer.background.celery.tasks.pruning",
-    ]
-)

backend/danswer/background/celery/apps/indexing.py

@@ -1,88 +0,0 @@
-import multiprocessing
-from typing import Any
-
-from celery import Celery
-from celery import signals
-from celery import Task
-from celery.signals import celeryd_init
-from celery.signals import worker_init
-from celery.signals import worker_ready
-from celery.signals import worker_shutdown
-
-import danswer.background.celery.apps.app_base as app_base
-from danswer.configs.constants import POSTGRES_CELERY_WORKER_INDEXING_APP_NAME
-from danswer.db.engine import SqlEngine
-from danswer.utils.logger import setup_logger
-
-
-logger = setup_logger()
-
-celery_app = Celery(__name__)
-celery_app.config_from_object("danswer.background.celery.configs.indexing")
-
-
-@signals.task_prerun.connect
-def on_task_prerun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)
-
-
-@signals.task_postrun.connect
-def on_task_postrun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    retval: Any | None = None,
-    state: str | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)
-
-
-@celeryd_init.connect
-def on_celeryd_init(sender: Any = None, conf: Any = None, **kwargs: Any) -> None:
-    app_base.on_celeryd_init(sender, conf, **kwargs)
-
-
-@worker_init.connect
-def on_worker_init(sender: Any, **kwargs: Any) -> None:
-    logger.info("worker_init signal received.")
-    logger.info(f"Multiprocessing start method: {multiprocessing.get_start_method()}")
-
-    SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_APP_NAME)
-    SqlEngine.init_engine(pool_size=8, max_overflow=0)
-
-    app_base.wait_for_redis(sender, **kwargs)
-    app_base.on_secondary_worker_init(sender, **kwargs)
-
-
-@worker_ready.connect
-def on_worker_ready(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_ready(sender, **kwargs)
-
-
-@worker_shutdown.connect
-def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_shutdown(sender, **kwargs)
-
-
-@signals.setup_logging.connect
-def on_setup_logging(
-    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
-) -> None:
-    app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
-
-
-celery_app.autodiscover_tasks(
-    [
-        "danswer.background.celery.tasks.indexing",
-    ]
-)

backend/danswer/background/celery/apps/light.py

@@ -1,89 +0,0 @@
-import multiprocessing
-from typing import Any
-
-from celery import Celery
-from celery import signals
-from celery import Task
-from celery.signals import celeryd_init
-from celery.signals import worker_init
-from celery.signals import worker_ready
-from celery.signals import worker_shutdown
-
-import danswer.background.celery.apps.app_base as app_base
-from danswer.configs.constants import POSTGRES_CELERY_WORKER_LIGHT_APP_NAME
-from danswer.db.engine import SqlEngine
-from danswer.utils.logger import setup_logger
-
-
-logger = setup_logger()
-
-celery_app = Celery(__name__)
-celery_app.config_from_object("danswer.background.celery.configs.light")
-
-
-@signals.task_prerun.connect
-def on_task_prerun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)
-
-
-@signals.task_postrun.connect
-def on_task_postrun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    retval: Any | None = None,
-    state: str | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)
-
-
-@celeryd_init.connect
-def on_celeryd_init(sender: Any = None, conf: Any = None, **kwargs: Any) -> None:
-    app_base.on_celeryd_init(sender, conf, **kwargs)
-
-
-@worker_init.connect
-def on_worker_init(sender: Any, **kwargs: Any) -> None:
-    logger.info("worker_init signal received.")
-    logger.info(f"Multiprocessing start method: {multiprocessing.get_start_method()}")
-
-    SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_LIGHT_APP_NAME)
-    SqlEngine.init_engine(pool_size=sender.concurrency, max_overflow=8)
-
-    app_base.wait_for_redis(sender, **kwargs)
-    app_base.on_secondary_worker_init(sender, **kwargs)
-
-
-@worker_ready.connect
-def on_worker_ready(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_ready(sender, **kwargs)
-
-
-@worker_shutdown.connect
-def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_shutdown(sender, **kwargs)
-
-
-@signals.setup_logging.connect
-def on_setup_logging(
-    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
-) -> None:
-    app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
-
-
-celery_app.autodiscover_tasks(
-    [
-        "danswer.background.celery.tasks.shared",
-        "danswer.background.celery.tasks.vespa",
-    ]
-)

backend/danswer/background/celery/apps/primary.py

@@ -1,300 +0,0 @@
-import multiprocessing
-from typing import Any
-
-from celery import bootsteps  # type: ignore
-from celery import Celery
-from celery import signals
-from celery import Task
-from celery.exceptions import WorkerShutdown
-from celery.signals import celeryd_init
-from celery.signals import worker_init
-from celery.signals import worker_ready
-from celery.signals import worker_shutdown
-
-import danswer.background.celery.apps.app_base as app_base
-from danswer.background.celery.apps.app_base import task_logger
-from danswer.background.celery.celery_redis import RedisConnectorCredentialPair
-from danswer.background.celery.celery_redis import RedisConnectorDeletion
-from danswer.background.celery.celery_redis import RedisConnectorIndexing
-from danswer.background.celery.celery_redis import RedisConnectorPruning
-from danswer.background.celery.celery_redis import RedisConnectorStop
-from danswer.background.celery.celery_redis import RedisDocumentSet
-from danswer.background.celery.celery_redis import RedisUserGroup
-from danswer.background.celery.celery_utils import celery_is_worker_primary
-from danswer.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT
-from danswer.configs.constants import DanswerRedisLocks
-from danswer.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME
-from danswer.db.engine import get_all_tenant_ids
-from danswer.db.engine import SqlEngine
-from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import setup_logger
-
-
-logger = setup_logger()
-
-celery_app = Celery(__name__)
-celery_app.config_from_object("danswer.background.celery.configs.primary")
-
-
-@signals.task_prerun.connect
-def on_task_prerun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_prerun(sender, task_id, task, args, kwargs, **kwds)
-
-
-@signals.task_postrun.connect
-def on_task_postrun(
-    sender: Any | None = None,
-    task_id: str | None = None,
-    task: Task | None = None,
-    args: tuple | None = None,
-    kwargs: dict | None = None,
-    retval: Any | None = None,
-    state: str | None = None,
-    **kwds: Any,
-) -> None:
-    app_base.on_task_postrun(sender, task_id, task, args, kwargs, retval, state, **kwds)
-
-
-@celeryd_init.connect
-def on_celeryd_init(sender: Any = None, conf: Any = None, **kwargs: Any) -> None:
-    app_base.on_celeryd_init(sender, conf, **kwargs)
-
-
-@worker_init.connect
-def on_worker_init(sender: Any, **kwargs: Any) -> None:
-    logger.info("worker_init signal received.")
-    logger.info(f"Multiprocessing start method: {multiprocessing.get_start_method()}")
-
-    SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME)
-    SqlEngine.init_engine(pool_size=8, max_overflow=0)
-
-    app_base.wait_for_redis(sender, **kwargs)
-
-    logger.info("Running as the primary celery worker.")
-
-    sender.primary_worker_locks = {}
-
-    # This is singleton work that should be done on startup exactly once
-    # by the primary worker
-    tenant_ids = get_all_tenant_ids()
-    for tenant_id in tenant_ids:
-        r = get_redis_client(tenant_id=tenant_id)
-
-        # For the moment, we're assuming that we are the only primary worker
-        # that should be running.
-        # TODO: maybe check for or clean up another zombie primary worker if we detect it
-        r.delete(DanswerRedisLocks.PRIMARY_WORKER)
-
-        # this process wide lock is taken to help other workers start up in order.
-        # it is planned to use this lock to enforce singleton behavior on the primary
-        # worker, since the primary worker does redis cleanup on startup, but this isn't
-        # implemented yet.
-        lock = r.lock(
-            DanswerRedisLocks.PRIMARY_WORKER,
-            timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
-        )
-
-        logger.info("Primary worker lock: Acquire starting.")
-        acquired = lock.acquire(blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2)
-        if acquired:
-            logger.info("Primary worker lock: Acquire succeeded.")
-        else:
-            logger.error("Primary worker lock: Acquire failed!")
-            raise WorkerShutdown("Primary worker lock could not be acquired!")
-
-        # tacking on our own user data to the sender
-        sender.primary_worker_locks[tenant_id] = lock
-
-        # As currently designed, when this worker starts as "primary", we reinitialize redis
-        # to a clean state (for our purposes, anyway)
-        r.delete(DanswerRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK)
-        r.delete(DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK)
-
-        r.delete(RedisConnectorCredentialPair.get_taskset_key())
-        r.delete(RedisConnectorCredentialPair.get_fence_key())
-
-        for key in r.scan_iter(RedisDocumentSet.TASKSET_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisUserGroup.TASKSET_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorDeletion.TASKSET_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorDeletion.FENCE_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorPruning.TASKSET_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorPruning.GENERATOR_COMPLETE_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorPruning.GENERATOR_PROGRESS_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorIndexing.TASKSET_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorIndexing.GENERATOR_COMPLETE_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorIndexing.GENERATOR_PROGRESS_PREFIX + "*"):
-            r.delete(key)
-
-        for key in r.scan_iter(RedisConnectorIndexing.FENCE_PREFIX + "*"):
-            r.delete(key)
-
-    for key in r.scan_iter(RedisConnectorStop.FENCE_PREFIX + "*"):
-        r.delete(key)
-
-
-# @worker_process_init.connect
-# def on_worker_process_init(sender: Any, **kwargs: Any) -> None:
-#     """This only runs inside child processes when the worker is in pool=prefork mode.
-#     This may be technically unnecessary since we're finding prefork pools to be
-#     unstable and currently aren't planning on using them."""
-#     logger.info("worker_process_init signal received.")
-#     SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME)
-#     SqlEngine.init_engine(pool_size=5, max_overflow=0)
-
-#     # https://stackoverflow.com/questions/43944787/sqlalchemy-celery-with-scoped-session-error
-#     SqlEngine.get_engine().dispose(close=False)
-
-
-@worker_ready.connect
-def on_worker_ready(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_ready(sender, **kwargs)
-
-
-@worker_shutdown.connect
-def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
-    app_base.on_worker_shutdown(sender, **kwargs)
-
-
-@signals.setup_logging.connect
-def on_setup_logging(
-    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
-) -> None:
-    app_base.on_setup_logging(loglevel, logfile, format, colorize, **kwargs)
-
-
-class HubPeriodicTask(bootsteps.StartStopStep):
-    """Regularly reacquires the primary worker lock outside of the task queue.
-    Use the task_logger in this class to avoid double logging.
-
-    This cannot be done inside a regular beat task because it must run on schedule and
-    a queue of existing work would starve the task from running.
-    """
-
-    # it's unclear to me whether using the hub's timer or the bootstep timer is better
-    requires = {"celery.worker.components:Hub"}
-
-    def __init__(self, worker: Any, **kwargs: Any) -> None:
-        self.interval = CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 8  # Interval in seconds
-        self.task_tref = None
-
-    def start(self, worker: Any) -> None:
-        if not celery_is_worker_primary(worker):
-            return
-
-        # Access the worker's event loop (hub)
-        hub = worker.consumer.controller.hub
-
-        # Schedule the periodic task
-        self.task_tref = hub.call_repeatedly(
-            self.interval, self.run_periodic_task, worker
-        )
-        task_logger.info("Scheduled periodic task with hub.")
-
-    def run_periodic_task(self, worker: Any) -> None:
-        try:
-            if not celery_is_worker_primary(worker):
-                return
-
-            if not hasattr(worker, "primary_worker_locks"):
-                return
-
-            # Retrieve all tenant IDs
-            tenant_ids = get_all_tenant_ids()
-
-            for tenant_id in tenant_ids:
-                lock = worker.primary_worker_locks.get(tenant_id)
-                if not lock:
-                    continue  # Skip if no lock for this tenant
-
-                r = get_redis_client(tenant_id=tenant_id)
-
-                if lock.owned():
-                    task_logger.debug(
-                        f"Reacquiring primary worker lock for tenant {tenant_id}."
-                    )
-                    lock.reacquire()
-                else:
-                    task_logger.warning(
-                        f"Full acquisition of primary worker lock for tenant {tenant_id}. "
-                        "Reasons could be worker restart or lock expiration."
-                    )
-                    lock = r.lock(
-                        DanswerRedisLocks.PRIMARY_WORKER,
-                        timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
-                    )
-
-                    task_logger.info(
-                        f"Primary worker lock for tenant {tenant_id}: Acquire starting."
-                    )
-                    acquired = lock.acquire(
-                        blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2
-                    )
-                    if acquired:
-                        task_logger.info(
-                            f"Primary worker lock for tenant {tenant_id}: Acquire succeeded."
-                        )
-                        worker.primary_worker_locks[tenant_id] = lock
-                    else:
-                        task_logger.error(
-                            f"Primary worker lock for tenant {tenant_id}: Acquire failed!"
-                        )
-                        raise TimeoutError(
-                            f"Primary worker lock for tenant {tenant_id} could not be acquired!"
-                        )
-
-        except Exception:
-            task_logger.exception("Periodic task failed.")
-
-    def stop(self, worker: Any) -> None:
-        # Cancel the scheduled task when the worker stops
-        if self.task_tref:
-            self.task_tref.cancel()
-            task_logger.info("Canceled periodic task with hub.")
-
-
-celery_app.steps["worker"].add(HubPeriodicTask)
-
-celery_app.autodiscover_tasks(
-    [
-        "danswer.background.celery.tasks.connector_deletion",
-        "danswer.background.celery.tasks.indexing",
-        "danswer.background.celery.tasks.periodic",
-        "danswer.background.celery.tasks.pruning",
-        "danswer.background.celery.tasks.shared",
-        "danswer.background.celery.tasks.vespa",
-    ]
-)

backend/danswer/background/celery/apps/task_formatters.py

@@ -1,26 +0,0 @@
-import logging
-
-from celery import current_task
-
-from danswer.utils.logger import ColoredFormatter
-from danswer.utils.logger import PlainFormatter
-
-
-class CeleryTaskPlainFormatter(PlainFormatter):
-    def format(self, record: logging.LogRecord) -> str:
-        task = current_task
-        if task and task.request:
-            record.__dict__.update(task_id=task.request.id, task_name=task.name)
-            record.msg = f"[{task.name}({task.request.id})] {record.msg}"
-
-        return super().format(record)
-
-
-class CeleryTaskColoredFormatter(ColoredFormatter):
-    def format(self, record: logging.LogRecord) -> str:
-        task = current_task
-        if task and task.request:
-            record.__dict__.update(task_id=task.request.id, task_name=task.name)
-            record.msg = f"[{task.name}({task.request.id})] {record.msg}"
-
-        return super().format(record)

backend/danswer/background/celery/celery_app.py

@@ -0,0 +1,484 @@
+import logging
+import time
+from datetime import timedelta
+from typing import Any
+
+import redis
+from celery import bootsteps  # type: ignore
+from celery import Celery
+from celery import current_task
+from celery import signals
+from celery import Task
+from celery.exceptions import WorkerShutdown
+from celery.signals import beat_init
+from celery.signals import worker_init
+from celery.signals import worker_ready
+from celery.signals import worker_shutdown
+from celery.states import READY_STATES
+from celery.utils.log import get_task_logger
+
+from danswer.background.celery.celery_redis import RedisConnectorCredentialPair
+from danswer.background.celery.celery_redis import RedisConnectorDeletion
+from danswer.background.celery.celery_redis import RedisConnectorPruning
+from danswer.background.celery.celery_redis import RedisDocumentSet
+from danswer.background.celery.celery_redis import RedisUserGroup
+from danswer.background.celery.celery_utils import celery_is_worker_primary
+from danswer.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT
+from danswer.configs.constants import DanswerCeleryPriority
+from danswer.configs.constants import DanswerRedisLocks
+from danswer.configs.constants import POSTGRES_CELERY_BEAT_APP_NAME
+from danswer.configs.constants import POSTGRES_CELERY_WORKER_HEAVY_APP_NAME
+from danswer.configs.constants import POSTGRES_CELERY_WORKER_LIGHT_APP_NAME
+from danswer.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME
+from danswer.db.engine import SqlEngine
+from danswer.redis.redis_pool import get_redis_client
+from danswer.utils.logger import ColoredFormatter
+from danswer.utils.logger import PlainFormatter
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+# use this within celery tasks to get celery task specific logging
+task_logger = get_task_logger(__name__)
+
+celery_app = Celery(__name__)
+celery_app.config_from_object(
+    "danswer.background.celery.celeryconfig"
+)  # Load configuration from 'celeryconfig.py'
+
+
+@signals.task_postrun.connect
+def celery_task_postrun(
+    sender: Any | None = None,
+    task_id: str | None = None,
+    task: Task | None = None,
+    args: tuple | None = None,
+    kwargs: dict | None = None,
+    retval: Any | None = None,
+    state: str | None = None,
+    **kwds: Any,
+) -> None:
+    """We handle this signal in order to remove completed tasks
+    from their respective tasksets. This allows us to track the progress of document set
+    and user group syncs.
+
+    This function runs after any task completes (both success and failure)
+    Note that this signal does not fire on a task that failed to complete and is going
+    to be retried.
+    """
+    if not task:
+        return
+
+    task_logger.debug(f"Task {task.name} (ID: {task_id}) completed with state: {state}")
+    # logger.debug(f"Result: {retval}")
+
+    if state not in READY_STATES:
+        return
+
+    if not task_id:
+        return
+
+    r = get_redis_client()
+
+    if task_id.startswith(RedisConnectorCredentialPair.PREFIX):
+        r.srem(RedisConnectorCredentialPair.get_taskset_key(), task_id)
+        return
+
+    if task_id.startswith(RedisDocumentSet.PREFIX):
+        document_set_id = RedisDocumentSet.get_id_from_task_id(task_id)
+        if document_set_id is not None:
+            rds = RedisDocumentSet(document_set_id)
+            r.srem(rds.taskset_key, task_id)
+        return
+
+    if task_id.startswith(RedisUserGroup.PREFIX):
+        usergroup_id = RedisUserGroup.get_id_from_task_id(task_id)
+        if usergroup_id is not None:
+            rug = RedisUserGroup(usergroup_id)
+            r.srem(rug.taskset_key, task_id)
+        return
+
+    if task_id.startswith(RedisConnectorDeletion.PREFIX):
+        cc_pair_id = RedisConnectorDeletion.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            rcd = RedisConnectorDeletion(cc_pair_id)
+            r.srem(rcd.taskset_key, task_id)
+        return
+
+    if task_id.startswith(RedisConnectorPruning.SUBTASK_PREFIX):
+        cc_pair_id = RedisConnectorPruning.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            rcp = RedisConnectorPruning(cc_pair_id)
+            r.srem(rcp.taskset_key, task_id)
+        return
+
+
+@beat_init.connect
+def on_beat_init(sender: Any, **kwargs: Any) -> None:
+    SqlEngine.set_app_name(POSTGRES_CELERY_BEAT_APP_NAME)
+    SqlEngine.init_engine(pool_size=2, max_overflow=0)
+
+
+@worker_init.connect
+def on_worker_init(sender: Any, **kwargs: Any) -> None:
+    # decide some initial startup settings based on the celery worker's hostname
+    # (set at the command line)
+    hostname = sender.hostname
+    if hostname.startswith("light"):
+        SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_LIGHT_APP_NAME)
+        SqlEngine.init_engine(pool_size=sender.concurrency, max_overflow=8)
+    elif hostname.startswith("heavy"):
+        SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_HEAVY_APP_NAME)
+        SqlEngine.init_engine(pool_size=8, max_overflow=0)
+    else:
+        SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME)
+        SqlEngine.init_engine(pool_size=8, max_overflow=0)
+
+    r = get_redis_client()
+
+    WAIT_INTERVAL = 5
+    WAIT_LIMIT = 60
+
+    time_start = time.monotonic()
+    logger.info("Redis: Readiness check starting.")
+    while True:
+        try:
+            if r.ping():
+                break
+        except Exception:
+            pass
+
+        time_elapsed = time.monotonic() - time_start
+        logger.info(
+            f"Redis: Ping failed. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
+        )
+        if time_elapsed > WAIT_LIMIT:
+            msg = (
+                f"Redis: Readiness check did not succeed within the timeout "
+                f"({WAIT_LIMIT} seconds). Exiting..."
+            )
+            logger.error(msg)
+            raise WorkerShutdown(msg)
+
+        time.sleep(WAIT_INTERVAL)
+
+    logger.info("Redis: Readiness check succeeded. Continuing...")
+
+    if not celery_is_worker_primary(sender):
+        logger.info("Running as a secondary celery worker.")
+        logger.info("Waiting for primary worker to be ready...")
+        time_start = time.monotonic()
+        while True:
+            if r.exists(DanswerRedisLocks.PRIMARY_WORKER):
+                break
+
+            time.monotonic()
+            time_elapsed = time.monotonic() - time_start
+            logger.info(
+                f"Primary worker is not ready yet. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
+            )
+            if time_elapsed > WAIT_LIMIT:
+                msg = (
+                    f"Primary worker was not ready within the timeout. "
+                    f"({WAIT_LIMIT} seconds). Exiting..."
+                )
+                logger.error(msg)
+                raise WorkerShutdown(msg)
+
+            time.sleep(WAIT_INTERVAL)
+
+        logger.info("Wait for primary worker completed successfully. Continuing...")
+        return
+
+    logger.info("Running as the primary celery worker.")
+
+    # This is singleton work that should be done on startup exactly once
+    # by the primary worker
+    r = get_redis_client()
+
+    # For the moment, we're assuming that we are the only primary worker
+    # that should be running.
+    # TODO: maybe check for or clean up another zombie primary worker if we detect it
+    r.delete(DanswerRedisLocks.PRIMARY_WORKER)
+
+    # this process wide lock is taken to help other workers start up in order.
+    # it is planned to use this lock to enforce singleton behavior on the primary
+    # worker, since the primary worker does redis cleanup on startup, but this isn't
+    # implemented yet.
+    lock = r.lock(
+        DanswerRedisLocks.PRIMARY_WORKER,
+        timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
+    )
+
+    logger.info("Primary worker lock: Acquire starting.")
+    acquired = lock.acquire(blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2)
+    if acquired:
+        logger.info("Primary worker lock: Acquire succeeded.")
+    else:
+        logger.error("Primary worker lock: Acquire failed!")
+        raise WorkerShutdown("Primary worker lock could not be acquired!")
+
+    sender.primary_worker_lock = lock
+
+    r.delete(DanswerRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK)
+    r.delete(DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK)
+
+    r.delete(RedisConnectorCredentialPair.get_taskset_key())
+    r.delete(RedisConnectorCredentialPair.get_fence_key())
+
+    for key in r.scan_iter(RedisDocumentSet.TASKSET_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisUserGroup.TASKSET_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisConnectorDeletion.TASKSET_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisConnectorDeletion.FENCE_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisConnectorPruning.TASKSET_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisConnectorPruning.GENERATOR_COMPLETE_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisConnectorPruning.GENERATOR_PROGRESS_PREFIX + "*"):
+        r.delete(key)
+
+    for key in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
+        r.delete(key)
+
+
+@worker_ready.connect
+def on_worker_ready(sender: Any, **kwargs: Any) -> None:
+    task_logger.info("worker_ready signal received.")
+
+
+@worker_shutdown.connect
+def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
+    if not celery_is_worker_primary(sender):
+        return
+
+    if not sender.primary_worker_lock:
+        return
+
+    logger.info("Releasing primary worker lock.")
+    lock = sender.primary_worker_lock
+    if lock.owned():
+        lock.release()
+        sender.primary_worker_lock = None
+
+
+class CeleryTaskPlainFormatter(PlainFormatter):
+    def format(self, record: logging.LogRecord) -> str:
+        task = current_task
+        if task and task.request:
+            record.__dict__.update(task_id=task.request.id, task_name=task.name)
+            record.msg = f"[{task.name}({task.request.id})] {record.msg}"
+
+        return super().format(record)
+
+
+class CeleryTaskColoredFormatter(ColoredFormatter):
+    def format(self, record: logging.LogRecord) -> str:
+        task = current_task
+        if task and task.request:
+            record.__dict__.update(task_id=task.request.id, task_name=task.name)
+            record.msg = f"[{task.name}({task.request.id})] {record.msg}"
+
+        return super().format(record)
+
+
+@signals.setup_logging.connect
+def on_setup_logging(
+    loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any
+) -> None:
+    # TODO: could unhardcode format and colorize and accept these as options from
+    # celery's config
+
+    # reformats celery's worker logger
+    root_logger = logging.getLogger()
+
+    root_handler = logging.StreamHandler()  # Set up a handler for the root logger
+    root_formatter = ColoredFormatter(
+        "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
+        datefmt="%m/%d/%Y %I:%M:%S %p",
+    )
+    root_handler.setFormatter(root_formatter)
+    root_logger.addHandler(root_handler)  # Apply the handler to the root logger
+
+    if logfile:
+        root_file_handler = logging.FileHandler(logfile)
+        root_file_formatter = PlainFormatter(
+            "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
+            datefmt="%m/%d/%Y %I:%M:%S %p",
+        )
+        root_file_handler.setFormatter(root_file_formatter)
+        root_logger.addHandler(root_file_handler)
+
+    root_logger.setLevel(loglevel)
+
+    # reformats celery's task logger
+    task_formatter = CeleryTaskColoredFormatter(
+        "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
+        datefmt="%m/%d/%Y %I:%M:%S %p",
+    )
+    task_handler = logging.StreamHandler()  # Set up a handler for the task logger
+    task_handler.setFormatter(task_formatter)
+    task_logger.addHandler(task_handler)  # Apply the handler to the task logger
+
+    if logfile:
+        task_file_handler = logging.FileHandler(logfile)
+        task_file_formatter = CeleryTaskPlainFormatter(
+            "%(asctime)s %(filename)30s %(lineno)4s: %(message)s",
+            datefmt="%m/%d/%Y %I:%M:%S %p",
+        )
+        task_file_handler.setFormatter(task_file_formatter)
+        task_logger.addHandler(task_file_handler)
+
+    task_logger.setLevel(loglevel)
+    task_logger.propagate = False
+
+
+class HubPeriodicTask(bootsteps.StartStopStep):
+    """Regularly reacquires the primary worker lock outside of the task queue.
+    Use the task_logger in this class to avoid double logging.
+
+    This cannot be done inside a regular beat task because it must run on schedule and
+    a queue of existing work would starve the task from running.
+    """
+
+    # it's unclear to me whether using the hub's timer or the bootstep timer is better
+    requires = {"celery.worker.components:Hub"}
+
+    def __init__(self, worker: Any, **kwargs: Any) -> None:
+        self.interval = CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 8  # Interval in seconds
+        self.task_tref = None
+
+    def start(self, worker: Any) -> None:
+        if not celery_is_worker_primary(worker):
+            return
+
+        # Access the worker's event loop (hub)
+        hub = worker.consumer.controller.hub
+
+        # Schedule the periodic task
+        self.task_tref = hub.call_repeatedly(
+            self.interval, self.run_periodic_task, worker
+        )
+        task_logger.info("Scheduled periodic task with hub.")
+
+    def run_periodic_task(self, worker: Any) -> None:
+        try:
+            if not worker.primary_worker_lock:
+                return
+
+            if not hasattr(worker, "primary_worker_lock"):
+                return
+
+            r = get_redis_client()
+
+            lock: redis.lock.Lock = worker.primary_worker_lock
+
+            if lock.owned():
+                task_logger.debug("Reacquiring primary worker lock.")
+                lock.reacquire()
+            else:
+                task_logger.warning(
+                    "Full acquisition of primary worker lock. "
+                    "Reasons could be computer sleep or a clock change."
+                )
+                lock = r.lock(
+                    DanswerRedisLocks.PRIMARY_WORKER,
+                    timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
+                )
+
+                task_logger.info("Primary worker lock: Acquire starting.")
+                acquired = lock.acquire(
+                    blocking_timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT / 2
+                )
+                if acquired:
+                    task_logger.info("Primary worker lock: Acquire succeeded.")
+                else:
+                    task_logger.error("Primary worker lock: Acquire failed!")
+                    raise TimeoutError("Primary worker lock could not be acquired!")
+
+                worker.primary_worker_lock = lock
+        except Exception:
+            task_logger.exception("HubPeriodicTask.run_periodic_task exceptioned.")
+
+    def stop(self, worker: Any) -> None:
+        # Cancel the scheduled task when the worker stops
+        if self.task_tref:
+            self.task_tref.cancel()
+            task_logger.info("Canceled periodic task with hub.")
+
+
+celery_app.steps["worker"].add(HubPeriodicTask)
+
+celery_app.autodiscover_tasks(
+    [
+        "danswer.background.celery.tasks.connector_deletion",
+        "danswer.background.celery.tasks.periodic",
+        "danswer.background.celery.tasks.pruning",
+        "danswer.background.celery.tasks.shared",
+        "danswer.background.celery.tasks.vespa",
+    ]
+)
+
+#####
+# Celery Beat (Periodic Tasks) Settings
+#####
+celery_app.conf.beat_schedule = {
+    "check-for-vespa-sync": {
+        "task": "check_for_vespa_sync_task",
+        "schedule": timedelta(seconds=5),
+        "options": {"priority": DanswerCeleryPriority.HIGH},
+    },
+}
+celery_app.conf.beat_schedule.update(
+    {
+        "check-for-connector-deletion-task": {
+            "task": "check_for_connector_deletion_task",
+            # don't need to check too often, since we kick off a deletion initially
+            # during the API call that actually marks the CC pair for deletion
+            "schedule": timedelta(seconds=60),
+            "options": {"priority": DanswerCeleryPriority.HIGH},
+        },
+    }
+)
+celery_app.conf.beat_schedule.update(
+    {
+        "check-for-prune": {
+            "task": "check_for_prune_task_2",
+            "schedule": timedelta(seconds=60),
+            "options": {"priority": DanswerCeleryPriority.HIGH},
+        },
+    }
+)
+celery_app.conf.beat_schedule.update(
+    {
+        "kombu-message-cleanup": {
+            "task": "kombu_message_cleanup_task",
+            "schedule": timedelta(seconds=3600),
+            "options": {"priority": DanswerCeleryPriority.LOWEST},
+        },
+    }
+)
+celery_app.conf.beat_schedule.update(
+    {
+        "monitor-vespa-sync": {
+            "task": "monitor_vespa_sync",
+            "schedule": timedelta(seconds=5),
+            "options": {"priority": DanswerCeleryPriority.HIGH},
+        },
+    }
+)

backend/danswer/background/celery/celery_redis.py

@@ -10,7 +10,7 @@ from celery import Celery
 from redis import Redis
 from sqlalchemy.orm import Session
 
-from danswer.background.celery.configs.base import CELERY_SEPARATOR
+from danswer.background.celery.celeryconfig import CELERY_SEPARATOR
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
 from danswer.configs.constants import DanswerCeleryPriority
 from danswer.configs.constants import DanswerCeleryQueues
@@ -21,7 +21,6 @@ from danswer.db.document import (
 )
 from danswer.db.document_set import construct_document_select_by_docset
 from danswer.utils.variable_functionality import fetch_versioned_implementation
-from danswer.utils.variable_functionality import global_version
 
 
 class RedisObjectHelper(ABC):
@@ -29,8 +28,8 @@ class RedisObjectHelper(ABC):
     FENCE_PREFIX = PREFIX + "_fence"
     TASKSET_PREFIX = PREFIX + "_taskset"
 
-    def __init__(self, id: str):
-        self._id: str = id
+    def __init__(self, id: int):
+        self._id: int = id
 
     @property
     def task_id_prefix(self) -> str:
@@ -47,7 +46,7 @@ class RedisObjectHelper(ABC):
         return f"{self.TASKSET_PREFIX}_{self._id}"
 
     @staticmethod
-    def get_id_from_fence_key(key: str) -> str | None:
+    def get_id_from_fence_key(key: str) -> int | None:
         """
         Extracts the object ID from a fence key in the format `PREFIX_fence_X`.
 
@@ -61,11 +60,15 @@ class RedisObjectHelper(ABC):
         if len(parts) != 3:
             return None
 
-        object_id = parts[2]
+        try:
+            object_id = int(parts[2])
+        except ValueError:
+            return None
+
         return object_id
 
     @staticmethod
-    def get_id_from_task_id(task_id: str) -> str | None:
+    def get_id_from_task_id(task_id: str) -> int | None:
         """
         Extracts the object ID from a task ID string.
 
@@ -89,7 +92,11 @@ class RedisObjectHelper(ABC):
         if len(parts) != 3:
             return None
 
-        object_id = parts[1]
+        try:
+            object_id = int(parts[1])
+        except ValueError:
+            return None
+
         return object_id
 
     @abstractmethod
@@ -99,7 +106,6 @@ class RedisObjectHelper(ABC):
         db_session: Session,
         redis_client: Redis,
         lock: redis.lock.Lock,
-        tenant_id: str | None,
     ) -> int | None:
         pass
 
@@ -109,21 +115,17 @@ class RedisDocumentSet(RedisObjectHelper):
     FENCE_PREFIX = PREFIX + "_fence"
     TASKSET_PREFIX = PREFIX + "_taskset"
 
-    def __init__(self, id: int) -> None:
-        super().__init__(str(id))
-
     def generate_tasks(
         self,
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
         lock: redis.lock.Lock,
-        tenant_id: str | None,
     ) -> int | None:
         last_lock_time = time.monotonic()
 
         async_results = []
-        stmt = construct_document_select_by_docset(int(self._id), current_only=False)
+        stmt = construct_document_select_by_docset(self._id, current_only=False)
         for doc in db_session.scalars(stmt).yield_per(1):
             current_time = time.monotonic()
             if current_time - last_lock_time >= (
@@ -143,7 +145,7 @@ class RedisDocumentSet(RedisObjectHelper):
 
             result = celery_app.send_task(
                 "vespa_metadata_sync_task",
-                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
+                kwargs=dict(document_id=doc.id),
                 queue=DanswerCeleryQueues.VESPA_METADATA_SYNC,
                 task_id=custom_task_id,
                 priority=DanswerCeleryPriority.LOW,
@@ -159,24 +161,17 @@ class RedisUserGroup(RedisObjectHelper):
     FENCE_PREFIX = PREFIX + "_fence"
     TASKSET_PREFIX = PREFIX + "_taskset"
 
-    def __init__(self, id: int) -> None:
-        super().__init__(str(id))
-
     def generate_tasks(
         self,
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
         lock: redis.lock.Lock,
-        tenant_id: str | None,
     ) -> int | None:
         last_lock_time = time.monotonic()
 
         async_results = []
 
-        if not global_version.is_ee_version():
-            return 0
-
         try:
             construct_document_select_by_usergroup = fetch_versioned_implementation(
                 "danswer.db.user_group",
@@ -185,7 +180,7 @@ class RedisUserGroup(RedisObjectHelper):
         except ModuleNotFoundError:
             return 0
 
-        stmt = construct_document_select_by_usergroup(int(self._id))
+        stmt = construct_document_select_by_usergroup(self._id)
         for doc in db_session.scalars(stmt).yield_per(1):
             current_time = time.monotonic()
             if current_time - last_lock_time >= (
@@ -205,7 +200,7 @@ class RedisUserGroup(RedisObjectHelper):
 
             result = celery_app.send_task(
                 "vespa_metadata_sync_task",
-                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
+                kwargs=dict(document_id=doc.id),
                 queue=DanswerCeleryQueues.VESPA_METADATA_SYNC,
                 task_id=custom_task_id,
                 priority=DanswerCeleryPriority.LOW,
@@ -217,19 +212,13 @@ class RedisUserGroup(RedisObjectHelper):
 
 
 class RedisConnectorCredentialPair(RedisObjectHelper):
-    """This class is used to scan documents by cc_pair in the db and collect them into
-    a unified set for syncing.
-
-    It differs from the other redis helpers in that the taskset used spans
+    """This class differs from the default in that the taskset used spans
     all connectors and is not per connector."""
 
     PREFIX = "connectorsync"
     FENCE_PREFIX = PREFIX + "_fence"
     TASKSET_PREFIX = PREFIX + "_taskset"
 
-    def __init__(self, id: int) -> None:
-        super().__init__(str(id))
-
     @classmethod
     def get_fence_key(cls) -> str:
         return RedisConnectorCredentialPair.FENCE_PREFIX
@@ -251,12 +240,11 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
         db_session: Session,
         redis_client: Redis,
         lock: redis.lock.Lock,
-        tenant_id: str | None,
     ) -> int | None:
         last_lock_time = time.monotonic()
 
         async_results = []
-        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
+        cc_pair = get_connector_credential_pair_from_id(self._id, db_session)
         if not cc_pair:
             return None
 
@@ -286,7 +274,7 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
             # Priority on sync's triggered by new indexing should be medium
             result = celery_app.send_task(
                 "vespa_metadata_sync_task",
-                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
+                kwargs=dict(document_id=doc.id),
                 queue=DanswerCeleryQueues.VESPA_METADATA_SYNC,
                 task_id=custom_task_id,
                 priority=DanswerCeleryPriority.MEDIUM,
@@ -302,23 +290,17 @@ class RedisConnectorDeletion(RedisObjectHelper):
     FENCE_PREFIX = PREFIX + "_fence"
     TASKSET_PREFIX = PREFIX + "_taskset"
 
-    def __init__(self, id: int) -> None:
-        super().__init__(str(id))
-
     def generate_tasks(
         self,
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
         lock: redis.lock.Lock,
-        tenant_id: str | None,
     ) -> int | None:
-        """Returns None if the cc_pair doesn't exist.
-        Otherwise, returns an int with the number of generated tasks."""
         last_lock_time = time.monotonic()
 
         async_results = []
-        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
+        cc_pair = get_connector_credential_pair_from_id(self._id, db_session)
         if not cc_pair:
             return None
 
@@ -350,7 +332,6 @@ class RedisConnectorDeletion(RedisObjectHelper):
                     document_id=doc.id,
                     connector_id=cc_pair.connector_id,
                     credential_id=cc_pair.credential_id,
-                    tenant_id=tenant_id,
                 ),
                 queue=DanswerCeleryQueues.CONNECTOR_DELETION,
                 task_id=custom_task_id,
@@ -395,7 +376,9 @@ class RedisConnectorPruning(RedisObjectHelper):
     )  # a signal that the generator has finished
 
     def __init__(self, id: int) -> None:
-        super().__init__(str(id))
+        """id: the cc_pair_id of the connector credential pair"""
+
+        super().__init__(id)
         self.documents_to_prune: set[str] = set()
 
     @property
@@ -422,12 +405,11 @@ class RedisConnectorPruning(RedisObjectHelper):
         db_session: Session,
         redis_client: Redis,
         lock: redis.lock.Lock | None,
-        tenant_id: str | None,
     ) -> int | None:
         last_lock_time = time.monotonic()
 
         async_results = []
-        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
+        cc_pair = get_connector_credential_pair_from_id(self._id, db_session)
         if not cc_pair:
             return None
 
@@ -456,7 +438,6 @@ class RedisConnectorPruning(RedisObjectHelper):
                     document_id=doc_id,
                     connector_id=cc_pair.connector_id,
                     credential_id=cc_pair.credential_id,
-                    tenant_id=tenant_id,
                 ),
                 queue=DanswerCeleryQueues.CONNECTOR_DELETION,
                 task_id=custom_task_id,
@@ -467,104 +448,20 @@ class RedisConnectorPruning(RedisObjectHelper):
 
         return len(async_results)
 
-    def is_pruning(self, redis_client: Redis) -> bool:
+    def is_pruning(self, db_session: Session, redis_client: Redis) -> bool:
         """A single example of a helper method being refactored into the redis helper"""
+        cc_pair = get_connector_credential_pair_from_id(
+            cc_pair_id=self._id, db_session=db_session
+        )
+        if not cc_pair:
+            raise ValueError(f"cc_pair_id {self._id} does not exist.")
+
         if redis_client.exists(self.fence_key):
             return True
 
         return False
 
 
-class RedisConnectorIndexing(RedisObjectHelper):
-    """Celery will kick off a long running indexing task to crawl the connector and
-    find any new or updated docs docs, which will each then get a new sync task or be
-    indexed inline.
-
-    ID should be a concatenation of cc_pair_id and search_setting_id, delimited by "/".
-    e.g. "2/5"
-    """
-
-    PREFIX = "connectorindexing"
-    FENCE_PREFIX = PREFIX + "_fence"  # a fence for the entire indexing process
-    GENERATOR_TASK_PREFIX = PREFIX + "+generator"
-
-    TASKSET_PREFIX = PREFIX + "_taskset"  # stores a list of prune tasks id's
-    SUBTASK_PREFIX = PREFIX + "+sub"
-
-    GENERATOR_LOCK_PREFIX = "da_lock:indexing"
-    GENERATOR_PROGRESS_PREFIX = (
-        PREFIX + "_generator_progress"
-    )  # a signal that contains generator progress
-    GENERATOR_COMPLETE_PREFIX = (
-        PREFIX + "_generator_complete"
-    )  # a signal that the generator has finished
-
-    def __init__(self, cc_pair_id: int, search_settings_id: int) -> None:
-        super().__init__(f"{cc_pair_id}/{search_settings_id}")
-
-    @property
-    def generator_lock_key(self) -> str:
-        return f"{self.GENERATOR_LOCK_PREFIX}_{self._id}"
-
-    @property
-    def generator_task_id_prefix(self) -> str:
-        return f"{self.GENERATOR_TASK_PREFIX}_{self._id}"
-
-    @property
-    def generator_progress_key(self) -> str:
-        # example: connectorpruning_generator_progress_1
-        return f"{self.GENERATOR_PROGRESS_PREFIX}_{self._id}"
-
-    @property
-    def generator_complete_key(self) -> str:
-        # example: connectorpruning_generator_complete_1
-        return f"{self.GENERATOR_COMPLETE_PREFIX}_{self._id}"
-
-    @property
-    def subtask_id_prefix(self) -> str:
-        return f"{self.SUBTASK_PREFIX}_{self._id}"
-
-    def generate_tasks(
-        self,
-        celery_app: Celery,
-        db_session: Session,
-        redis_client: Redis,
-        lock: redis.lock.Lock | None,
-        tenant_id: str | None,
-    ) -> int | None:
-        return None
-
-    def is_indexing(self, redis_client: Redis) -> bool:
-        """A single example of a helper method being refactored into the redis helper"""
-        if redis_client.exists(self.fence_key):
-            return True
-
-        return False
-
-
-class RedisConnectorStop(RedisObjectHelper):
-    """Used to signal any running tasks for a connector to stop. We should refactor
-    connector related redis helpers into a single class.
-    """
-
-    PREFIX = "connectorstop"
-    FENCE_PREFIX = PREFIX + "_fence"  # a fence for the entire indexing process
-    TASKSET_PREFIX = PREFIX + "_taskset"  # stores a list of prune tasks id's
-
-    def __init__(self, id: int) -> None:
-        super().__init__(str(id))
-
-    def generate_tasks(
-        self,
-        celery_app: Celery,
-        db_session: Session,
-        redis_client: Redis,
-        lock: redis.lock.Lock | None,
-        tenant_id: str | None,
-    ) -> int | None:
-        return None
-
-
 def celery_get_queue_length(queue: str, r: Redis) -> int:
     """This is a redis specific way to get the length of a celery queue.
     It is priority aware and knows how to count across the multiple redis lists

backend/danswer/background/celery/celery_run.py

@@ -1,8 +1,9 @@
-"""Factory stub for running celery worker / celery beat."""
+"""Entry point for running celery worker / celery beat."""
 from danswer.utils.variable_functionality import fetch_versioned_implementation
 from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
 
+
 set_is_ee_based_on_env_variable()
-app = fetch_versioned_implementation(
-    "danswer.background.celery.apps.primary", "celery_app"
+celery_app = fetch_versioned_implementation(
+    "danswer.background.celery.celery_app", "celery_app"
 )

backend/danswer/background/celery/celery_utils.py

@@ -1,3 +1,4 @@
+from collections.abc import Callable
 from datetime import datetime
 from datetime import timezone
 from typing import Any
@@ -5,15 +6,14 @@ from typing import Any
 from sqlalchemy.orm import Session
 
 from danswer.background.celery.celery_redis import RedisConnectorDeletion
-from danswer.background.indexing.run_indexing import RunIndexingCallbackInterface
 from danswer.configs.app_configs import MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE
 from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
     rate_limit_builder,
 )
 from danswer.connectors.interfaces import BaseConnector
+from danswer.connectors.interfaces import IdConnector
 from danswer.connectors.interfaces import LoadConnector
 from danswer.connectors.interfaces import PollConnector
-from danswer.connectors.interfaces import SlimConnector
 from danswer.connectors.models import Document
 from danswer.db.connector_credential_pair import get_connector_credential_pair
 from danswer.db.enums import TaskStatus
@@ -27,10 +27,7 @@ logger = setup_logger()
 
 
 def _get_deletion_status(
-    connector_id: int,
-    credential_id: int,
-    db_session: Session,
-    tenant_id: str | None = None,
+    connector_id: int, credential_id: int, db_session: Session
 ) -> TaskQueueState | None:
     """We no longer store TaskQueueState in the DB for a deletion attempt.
     This function populates TaskQueueState by just checking redis.
@@ -43,7 +40,7 @@ def _get_deletion_status(
 
     rcd = RedisConnectorDeletion(cc_pair.id)
 
-    r = get_redis_client(tenant_id=tenant_id)
+    r = get_redis_client()
     if not r.exists(rcd.fence_key):
         return None
 
@@ -53,14 +50,9 @@ def _get_deletion_status(
 
 
 def get_deletion_attempt_snapshot(
-    connector_id: int,
-    credential_id: int,
-    db_session: Session,
-    tenant_id: str | None = None,
+    connector_id: int, credential_id: int, db_session: Session
 ) -> DeletionAttemptSnapshot | None:
-    deletion_task = _get_deletion_status(
-        connector_id, credential_id, db_session, tenant_id
-    )
+    deletion_task = _get_deletion_status(connector_id, credential_id, db_session)
     if not deletion_task:
         return None
 
@@ -71,15 +63,13 @@ def get_deletion_attempt_snapshot(
     )
 
 
-def document_batch_to_ids(
-    doc_batch: list[Document],
-) -> set[str]:
+def document_batch_to_ids(doc_batch: list[Document]) -> set[str]:
     return {doc.id for doc in doc_batch}
 
 
 def extract_ids_from_runnable_connector(
     runnable_connector: BaseConnector,
-    callback: RunIndexingCallbackInterface | None = None,
+    progress_callback: Callable[[int], None] | None = None,
 ) -> set[str]:
     """
     If the PruneConnector hasnt been implemented for the given connector, just pull
@@ -89,13 +79,10 @@ def extract_ids_from_runnable_connector(
     """
     all_connector_doc_ids: set[str] = set()
 
-    if isinstance(runnable_connector, SlimConnector):
-        for metadata_batch in runnable_connector.retrieve_all_slim_documents():
-            all_connector_doc_ids.update({doc.id for doc in metadata_batch})
-
     doc_batch_generator = None
-
-    if isinstance(runnable_connector, LoadConnector):
+    if isinstance(runnable_connector, IdConnector):
+        all_connector_doc_ids = runnable_connector.retrieve_all_source_ids()
+    elif isinstance(runnable_connector, LoadConnector):
         doc_batch_generator = runnable_connector.load_from_state()
     elif isinstance(runnable_connector, PollConnector):
         start = datetime(1970, 1, 1, tzinfo=timezone.utc).timestamp()
@@ -104,17 +91,16 @@ def extract_ids_from_runnable_connector(
     else:
         raise RuntimeError("Pruning job could not find a valid runnable_connector.")
 
-    doc_batch_processing_func = document_batch_to_ids
-    if MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE:
-        doc_batch_processing_func = rate_limit_builder(
-            max_calls=MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE, period=60
-        )(document_batch_to_ids)
-    for doc_batch in doc_batch_generator:
-        if callback:
-            if callback.should_stop():
-                raise RuntimeError("Stop signal received")
-            callback.progress(len(doc_batch))
-        all_connector_doc_ids.update(doc_batch_processing_func(doc_batch))
+    if doc_batch_generator:
+        doc_batch_processing_func = document_batch_to_ids
+        if MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE:
+            doc_batch_processing_func = rate_limit_builder(
+                max_calls=MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE, period=60
+            )(document_batch_to_ids)
+        for doc_batch in doc_batch_generator:
+            if progress_callback:
+                progress_callback(len(doc_batch))
+            all_connector_doc_ids.update(doc_batch_processing_func(doc_batch))
 
     return all_connector_doc_ids
 
@@ -135,10 +121,13 @@ def celery_is_listening_to_queue(worker: Any, name: str) -> bool:
 def celery_is_worker_primary(worker: Any) -> bool:
     """There are multiple approaches that could be taken to determine if a celery worker
     is 'primary', as defined by us. But the way we do it is to check the hostname set
-    for the celery worker, which can be done on the
+    for the celery worker, which can be done either in celeryconfig.py or on the
     command line with '--hostname'."""
     hostname = worker.hostname
-    if hostname.startswith("primary"):
-        return True
+    if hostname.startswith("light"):
+        return False
 
-    return False
+    if hostname.startswith("heavy"):
+        return False
+
+    return True

backend/danswer/background/celery/celeryconfig.py

@@ -1,6 +1,4 @@
 # docs: https://docs.celeryq.dev/en/stable/userguide/configuration.html
-import urllib.parse
-
 from danswer.configs.app_configs import CELERY_BROKER_POOL_LIMIT
 from danswer.configs.app_configs import CELERY_RESULT_EXPIRES
 from danswer.configs.app_configs import REDIS_DB_NUMBER_CELERY
@@ -19,7 +17,7 @@ CELERY_SEPARATOR = ":"
 
 CELERY_PASSWORD_PART = ""
 if REDIS_PASSWORD:
-    CELERY_PASSWORD_PART = ":" + urllib.parse.quote(REDIS_PASSWORD, safe="") + "@"
+    CELERY_PASSWORD_PART = f":{REDIS_PASSWORD}@"
 
 REDIS_SCHEME = "redis"
 
@@ -31,10 +29,16 @@ if REDIS_SSL:
     if REDIS_SSL_CA_CERTS:
         SSL_QUERY_PARAMS += f"&ssl_ca_certs={REDIS_SSL_CA_CERTS}"
 
-# region Broker settings
 # example celery_broker_url: "redis://:password@localhost:6379/15"
 broker_url = f"{REDIS_SCHEME}://{CELERY_PASSWORD_PART}{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB_NUMBER_CELERY}{SSL_QUERY_PARAMS}"
 
+result_backend = f"{REDIS_SCHEME}://{CELERY_PASSWORD_PART}{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB_NUMBER_CELERY_RESULT_BACKEND}{SSL_QUERY_PARAMS}"
+
+# NOTE: prefetch 4 is significantly faster than prefetch 1 for small tasks
+# however, prefetching is bad when tasks are lengthy as those tasks
+# can stall other tasks.
+worker_prefetch_multiplier = 4
+
 broker_connection_retry_on_startup = True
 broker_pool_limit = CELERY_BROKER_POOL_LIMIT
 
@@ -49,7 +53,6 @@ broker_transport_options = {
     "socket_keepalive": True,
     "socket_keepalive_options": REDIS_SOCKET_KEEPALIVE_OPTIONS,
 }
-# endregion
 
 # redis backend settings
 # https://docs.celeryq.dev/en/stable/userguide/configuration.html#redis-backend-settings
@@ -63,19 +66,10 @@ redis_backend_health_check_interval = REDIS_HEALTH_CHECK_INTERVAL
 task_default_priority = DanswerCeleryPriority.MEDIUM
 task_acks_late = True
 
-# region Task result backend settings
 # It's possible we don't even need celery's result backend, in which case all of the optimization below
 # might be irrelevant
-result_backend = f"{REDIS_SCHEME}://{CELERY_PASSWORD_PART}{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB_NUMBER_CELERY_RESULT_BACKEND}{SSL_QUERY_PARAMS}"
 result_expires = CELERY_RESULT_EXPIRES  # 86400 seconds is the default
-# endregion
 
-# Leaving this to the default of True may cause double logging since both our own app
-# and celery think they are controlling the logger.
-# TODO: Configure celery's logger entirely manually and set this to False
-# worker_hijack_root_logger = False
-
-# region Notes on serialization performance
 # Option 0: Defaults (json serializer, no compression)
 # about 1.5 KB per queued task. 1KB in queue, 400B for result, 100 as a child entry in generator result
 
@@ -101,4 +95,3 @@ result_expires = CELERY_RESULT_EXPIRES  # 86400 seconds is the default
 # task_serializer = "pickle-bzip2"
 # result_serializer = "pickle-bzip2"
 # accept_content=["pickle", "pickle-bzip2"]
-# endregion

backend/danswer/background/celery/configs/beat.py

@@ -1,14 +0,0 @@
-# docs: https://docs.celeryq.dev/en/stable/userguide/configuration.html
-import danswer.background.celery.configs.base as shared_config
-
-broker_url = shared_config.broker_url
-broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup
-broker_pool_limit = shared_config.broker_pool_limit
-broker_transport_options = shared_config.broker_transport_options
-
-redis_socket_keepalive = shared_config.redis_socket_keepalive
-redis_retry_on_timeout = shared_config.redis_retry_on_timeout
-redis_backend_health_check_interval = shared_config.redis_backend_health_check_interval
-
-result_backend = shared_config.result_backend
-result_expires = shared_config.result_expires  # 86400 seconds is the default

backend/danswer/background/celery/configs/heavy.py

@@ -1,20 +0,0 @@
-import danswer.background.celery.configs.base as shared_config
-
-broker_url = shared_config.broker_url
-broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup
-broker_pool_limit = shared_config.broker_pool_limit
-broker_transport_options = shared_config.broker_transport_options
-
-redis_socket_keepalive = shared_config.redis_socket_keepalive
-redis_retry_on_timeout = shared_config.redis_retry_on_timeout
-redis_backend_health_check_interval = shared_config.redis_backend_health_check_interval
-
-result_backend = shared_config.result_backend
-result_expires = shared_config.result_expires  # 86400 seconds is the default
-
-task_default_priority = shared_config.task_default_priority
-task_acks_late = shared_config.task_acks_late
-
-worker_concurrency = 4
-worker_pool = "threads"
-worker_prefetch_multiplier = 1

backend/danswer/background/celery/configs/indexing.py

@@ -1,21 +0,0 @@
-import danswer.background.celery.configs.base as shared_config
-from danswer.configs.app_configs import CELERY_WORKER_INDEXING_CONCURRENCY
-
-broker_url = shared_config.broker_url
-broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup
-broker_pool_limit = shared_config.broker_pool_limit
-broker_transport_options = shared_config.broker_transport_options
-
-redis_socket_keepalive = shared_config.redis_socket_keepalive
-redis_retry_on_timeout = shared_config.redis_retry_on_timeout
-redis_backend_health_check_interval = shared_config.redis_backend_health_check_interval
-
-result_backend = shared_config.result_backend
-result_expires = shared_config.result_expires  # 86400 seconds is the default
-
-task_default_priority = shared_config.task_default_priority
-task_acks_late = shared_config.task_acks_late
-
-worker_concurrency = CELERY_WORKER_INDEXING_CONCURRENCY
-worker_pool = "threads"
-worker_prefetch_multiplier = 1

backend/danswer/background/celery/configs/light.py

@@ -1,22 +0,0 @@
-import danswer.background.celery.configs.base as shared_config
-from danswer.configs.app_configs import CELERY_WORKER_LIGHT_CONCURRENCY
-from danswer.configs.app_configs import CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER
-
-broker_url = shared_config.broker_url
-broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup
-broker_pool_limit = shared_config.broker_pool_limit
-broker_transport_options = shared_config.broker_transport_options
-
-redis_socket_keepalive = shared_config.redis_socket_keepalive
-redis_retry_on_timeout = shared_config.redis_retry_on_timeout
-redis_backend_health_check_interval = shared_config.redis_backend_health_check_interval
-
-result_backend = shared_config.result_backend
-result_expires = shared_config.result_expires  # 86400 seconds is the default
-
-task_default_priority = shared_config.task_default_priority
-task_acks_late = shared_config.task_acks_late
-
-worker_concurrency = CELERY_WORKER_LIGHT_CONCURRENCY
-worker_pool = "threads"
-worker_prefetch_multiplier = CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER

backend/danswer/background/celery/configs/primary.py

@@ -1,20 +0,0 @@
-import danswer.background.celery.configs.base as shared_config
-
-broker_url = shared_config.broker_url
-broker_connection_retry_on_startup = shared_config.broker_connection_retry_on_startup
-broker_pool_limit = shared_config.broker_pool_limit
-broker_transport_options = shared_config.broker_transport_options
-
-redis_socket_keepalive = shared_config.redis_socket_keepalive
-redis_retry_on_timeout = shared_config.redis_retry_on_timeout
-redis_backend_health_check_interval = shared_config.redis_backend_health_check_interval
-
-result_backend = shared_config.result_backend
-result_expires = shared_config.result_expires  # 86400 seconds is the default
-
-task_default_priority = shared_config.task_default_priority
-task_acks_late = shared_config.task_acks_late
-
-worker_concurrency = 4
-worker_pool = "threads"
-worker_prefetch_multiplier = 1

backend/danswer/background/celery/tasks/connector_deletion/tasks.py

@@ -1,46 +1,30 @@
-from datetime import datetime
-from datetime import timezone
-
 import redis
-from celery import Celery
 from celery import shared_task
-from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from redis import Redis
 from sqlalchemy.orm import Session
+from sqlalchemy.orm.exc import ObjectDeletedError
 
-from danswer.background.celery.apps.app_base import task_logger
+from danswer.background.celery.celery_app import celery_app
+from danswer.background.celery.celery_app import task_logger
 from danswer.background.celery.celery_redis import RedisConnectorDeletion
-from danswer.background.celery.celery_redis import RedisConnectorIndexing
-from danswer.background.celery.celery_redis import RedisConnectorPruning
-from danswer.background.celery.celery_redis import RedisConnectorStop
-from danswer.background.celery.tasks.shared.RedisConnectorDeletionFenceData import (
-    RedisConnectorDeletionFenceData,
-)
 from danswer.configs.app_configs import JOB_TIMEOUT
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
 from danswer.configs.constants import DanswerRedisLocks
-from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
 from danswer.db.connector_credential_pair import get_connector_credential_pairs
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.enums import ConnectorCredentialPairStatus
-from danswer.db.search_settings import get_all_search_settings
+from danswer.db.models import ConnectorCredentialPair
 from danswer.redis.redis_pool import get_redis_client
 
 
-class TaskDependencyError(RuntimeError):
-    """Raised to the caller to indicate dependent tasks are running that would interfere
-    with connector deletion."""
-
-
 @shared_task(
     name="check_for_connector_deletion_task",
     soft_time_limit=JOB_TIMEOUT,
     trail=False,
-    bind=True,
 )
-def check_for_connector_deletion_task(self: Task, *, tenant_id: str | None) -> None:
-    r = get_redis_client(tenant_id=tenant_id)
+def check_for_connector_deletion_task() -> None:
+    r = get_redis_client()
 
     lock_beat = r.lock(
         DanswerRedisLocks.CHECK_CONNECTOR_DELETION_BEAT_LOCK,
@@ -52,30 +36,12 @@ def check_for_connector_deletion_task(self: Task, *, tenant_id: str | None) -> N
         if not lock_beat.acquire(blocking=False):
             return
 
-        # collect cc_pair_ids
-        cc_pair_ids: list[int] = []
-        with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             cc_pairs = get_connector_credential_pairs(db_session)
             for cc_pair in cc_pairs:
-                cc_pair_ids.append(cc_pair.id)
-
-        # try running cleanup on the cc_pair_ids
-        for cc_pair_id in cc_pair_ids:
-            with get_session_with_tenant(tenant_id) as db_session:
-                rcs = RedisConnectorStop(cc_pair_id)
-                try:
-                    try_generate_document_cc_pair_cleanup_tasks(
-                        self.app, cc_pair_id, db_session, r, lock_beat, tenant_id
-                    )
-                except TaskDependencyError as e:
-                    # this means we wanted to start deleting but dependent tasks were running
-                    # Leave a stop signal to clear indexing and pruning tasks more quickly
-                    task_logger.info(str(e))
-                    r.set(rcs.fence_key, cc_pair_id)
-                else:
-                    # clear the stop signal if it exists ... no longer needed
-                    r.delete(rcs.fence_key)
-
+                try_generate_document_cc_pair_cleanup_tasks(
+                    cc_pair, db_session, r, lock_beat
+                )
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
@@ -88,97 +54,57 @@ def check_for_connector_deletion_task(self: Task, *, tenant_id: str | None) -> N
 
 
 def try_generate_document_cc_pair_cleanup_tasks(
-    app: Celery,
-    cc_pair_id: int,
+    cc_pair: ConnectorCredentialPair,
     db_session: Session,
     r: Redis,
     lock_beat: redis.lock.Lock,
-    tenant_id: str | None,
 ) -> int | None:
     """Returns an int if syncing is needed. The int represents the number of sync tasks generated.
     Note that syncing can still be required even if the number of sync tasks generated is zero.
     Returns None if no syncing is required.
-
-    Will raise TaskDependencyError if dependent tasks such as indexing and pruning are
-    still running. In our case, the caller reacts by setting a stop signal in Redis to
-    exit those tasks as quickly as possible.
     """
 
     lock_beat.reacquire()
 
-    rcd = RedisConnectorDeletion(cc_pair_id)
+    rcd = RedisConnectorDeletion(cc_pair.id)
 
     # don't generate sync tasks if tasks are still pending
     if r.exists(rcd.fence_key):
         return None
 
-    # we need to load the state of the object inside the fence
+    # we need to refresh the state of the object inside the fence
     # to avoid a race condition with db.commit/fence deletion
     # at the end of this taskset
-    cc_pair = get_connector_credential_pair_from_id(cc_pair_id, db_session)
-    if not cc_pair:
+    try:
+        db_session.refresh(cc_pair)
+    except ObjectDeletedError:
         return None
 
     if cc_pair.status != ConnectorCredentialPairStatus.DELETING:
         return None
 
-    # set a basic fence to start
-    fence_value = RedisConnectorDeletionFenceData(
-        num_tasks=None,
-        submitted=datetime.now(timezone.utc),
+    # add tasks to celery and build up the task set to monitor in redis
+    r.delete(rcd.taskset_key)
+
+    # Add all documents that need to be updated into the queue
+    task_logger.info(
+        f"RedisConnectorDeletion.generate_tasks starting. cc_pair_id={cc_pair.id}"
     )
-    r.set(rcd.fence_key, fence_value.model_dump_json())
-
-    try:
-        # do not proceed if connector indexing or connector pruning are running
-        search_settings_list = get_all_search_settings(db_session)
-        for search_settings in search_settings_list:
-            rci = RedisConnectorIndexing(cc_pair_id, search_settings.id)
-            if r.get(rci.fence_key):
-                raise TaskDependencyError(
-                    f"Connector deletion - Delayed (indexing in progress): "
-                    f"cc_pair={cc_pair_id} "
-                    f"search_settings={search_settings.id}"
-                )
-
-        rcp = RedisConnectorPruning(cc_pair_id)
-        if r.get(rcp.fence_key):
-            raise TaskDependencyError(
-                f"Connector deletion - Delayed (pruning in progress): "
-                f"cc_pair={cc_pair_id}"
-            )
-
-        # add tasks to celery and build up the task set to monitor in redis
-        r.delete(rcd.taskset_key)
-
-        # Add all documents that need to be updated into the queue
-        task_logger.info(
-            f"RedisConnectorDeletion.generate_tasks starting. cc_pair={cc_pair_id}"
-        )
-        tasks_generated = rcd.generate_tasks(app, db_session, r, lock_beat, tenant_id)
-        if tasks_generated is None:
-            raise ValueError("RedisConnectorDeletion.generate_tasks returned None")
-    except TaskDependencyError:
-        r.delete(rcd.fence_key)
-        raise
-    except Exception:
-        task_logger.exception("Unexpected exception")
-        r.delete(rcd.fence_key)
+    tasks_generated = rcd.generate_tasks(celery_app, db_session, r, lock_beat)
+    if tasks_generated is None:
         return None
-    else:
-        # Currently we are allowing the sync to proceed with 0 tasks.
-        # It's possible for sets/groups to be generated initially with no entries
-        # and they still need to be marked as up to date.
-        # if tasks_generated == 0:
-        #     return 0
 
-        task_logger.info(
-            f"RedisConnectorDeletion.generate_tasks finished. "
-            f"cc_pair={cc_pair_id} tasks_generated={tasks_generated}"
-        )
+    # Currently we are allowing the sync to proceed with 0 tasks.
+    # It's possible for sets/groups to be generated initially with no entries
+    # and they still need to be marked as up to date.
+    # if tasks_generated == 0:
+    #     return 0
 
-        # set this only after all tasks have been added
-        fence_value.num_tasks = tasks_generated
-        r.set(rcd.fence_key, fence_value.model_dump_json())
+    task_logger.info(
+        f"RedisConnectorDeletion.generate_tasks finished. "
+        f"cc_pair_id={cc_pair.id} tasks_generated={tasks_generated}"
+    )
 
+    # set this only after all tasks have been added
+    r.set(rcd.fence_key, tasks_generated)
     return tasks_generated

backend/danswer/background/celery/tasks/indexing/tasks.py

@@ -1,583 +0,0 @@
-from datetime import datetime
-from datetime import timezone
-from http import HTTPStatus
-from time import sleep
-from typing import cast
-from uuid import uuid4
-
-import redis
-from celery import Celery
-from celery import shared_task
-from celery import Task
-from celery.exceptions import SoftTimeLimitExceeded
-from redis import Redis
-from sqlalchemy.orm import Session
-
-from danswer.background.celery.apps.app_base import task_logger
-from danswer.background.celery.celery_redis import RedisConnectorDeletion
-from danswer.background.celery.celery_redis import RedisConnectorIndexing
-from danswer.background.celery.celery_redis import RedisConnectorStop
-from danswer.background.celery.tasks.shared.RedisConnectorIndexingFenceData import (
-    RedisConnectorIndexingFenceData,
-)
-from danswer.background.indexing.job_client import SimpleJobClient
-from danswer.background.indexing.run_indexing import run_indexing_entrypoint
-from danswer.background.indexing.run_indexing import RunIndexingCallbackInterface
-from danswer.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP
-from danswer.configs.app_configs import MULTI_TENANT
-from danswer.configs.constants import CELERY_INDEXING_LOCK_TIMEOUT
-from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
-from danswer.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX
-from danswer.configs.constants import DanswerCeleryPriority
-from danswer.configs.constants import DanswerCeleryQueues
-from danswer.configs.constants import DanswerRedisLocks
-from danswer.configs.constants import DocumentSource
-from danswer.db.connector_credential_pair import fetch_connector_credential_pairs
-from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
-from danswer.db.engine import get_db_current_time
-from danswer.db.engine import get_session_with_tenant
-from danswer.db.enums import ConnectorCredentialPairStatus
-from danswer.db.enums import IndexingStatus
-from danswer.db.enums import IndexModelStatus
-from danswer.db.index_attempt import create_index_attempt
-from danswer.db.index_attempt import get_index_attempt
-from danswer.db.index_attempt import get_last_attempt_for_cc_pair
-from danswer.db.index_attempt import mark_attempt_failed
-from danswer.db.models import ConnectorCredentialPair
-from danswer.db.models import IndexAttempt
-from danswer.db.models import SearchSettings
-from danswer.db.search_settings import get_current_search_settings
-from danswer.db.search_settings import get_secondary_search_settings
-from danswer.db.swap_index import check_index_swap
-from danswer.natural_language_processing.search_nlp_models import EmbeddingModel
-from danswer.natural_language_processing.search_nlp_models import warm_up_bi_encoder
-from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import setup_logger
-from danswer.utils.variable_functionality import global_version
-from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
-from shared_configs.configs import INDEXING_MODEL_SERVER_PORT
-
-logger = setup_logger()
-
-
-class RunIndexingCallback(RunIndexingCallbackInterface):
-    def __init__(
-        self,
-        stop_key: str,
-        generator_progress_key: str,
-        redis_lock: redis.lock.Lock,
-        redis_client: Redis,
-    ):
-        super().__init__()
-        self.redis_lock: redis.lock.Lock = redis_lock
-        self.stop_key: str = stop_key
-        self.generator_progress_key: str = generator_progress_key
-        self.redis_client = redis_client
-
-    def should_stop(self) -> bool:
-        if self.redis_client.exists(self.stop_key):
-            return True
-        return False
-
-    def progress(self, amount: int) -> None:
-        self.redis_lock.reacquire()
-        self.redis_client.incrby(self.generator_progress_key, amount)
-
-
-@shared_task(
-    name="check_for_indexing",
-    soft_time_limit=300,
-    bind=True,
-)
-def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
-    tasks_created = 0
-
-    r = get_redis_client(tenant_id=tenant_id)
-
-    lock_beat = r.lock(
-        DanswerRedisLocks.CHECK_INDEXING_BEAT_LOCK,
-        timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
-    )
-
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return None
-
-        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
-            check_index_swap(db_session=db_session)
-            current_search_settings = get_current_search_settings(db_session)
-            # So that the first time users aren't surprised by really slow speed of first
-            # batch of documents indexed
-            if current_search_settings.provider_type is None and not MULTI_TENANT:
-                embedding_model = EmbeddingModel.from_db_model(
-                    search_settings=current_search_settings,
-                    server_host=INDEXING_MODEL_SERVER_HOST,
-                    server_port=INDEXING_MODEL_SERVER_PORT,
-                )
-                warm_up_bi_encoder(
-                    embedding_model=embedding_model,
-                )
-
-        cc_pair_ids: list[int] = []
-        with get_session_with_tenant(tenant_id) as db_session:
-            cc_pairs = fetch_connector_credential_pairs(db_session)
-            for cc_pair_entry in cc_pairs:
-                cc_pair_ids.append(cc_pair_entry.id)
-
-        for cc_pair_id in cc_pair_ids:
-            with get_session_with_tenant(tenant_id) as db_session:
-                # Get the primary search settings
-                primary_search_settings = get_current_search_settings(db_session)
-                search_settings = [primary_search_settings]
-
-                # Check for secondary search settings
-                secondary_search_settings = get_secondary_search_settings(db_session)
-                if secondary_search_settings is not None:
-                    # If secondary settings exist, add them to the list
-                    search_settings.append(secondary_search_settings)
-
-                for search_settings_instance in search_settings:
-                    rci = RedisConnectorIndexing(
-                        cc_pair_id, search_settings_instance.id
-                    )
-                    if r.exists(rci.fence_key):
-                        continue
-
-                    cc_pair = get_connector_credential_pair_from_id(
-                        cc_pair_id, db_session
-                    )
-                    if not cc_pair:
-                        continue
-
-                    last_attempt = get_last_attempt_for_cc_pair(
-                        cc_pair.id, search_settings_instance.id, db_session
-                    )
-                    if not _should_index(
-                        cc_pair=cc_pair,
-                        last_index=last_attempt,
-                        search_settings_instance=search_settings_instance,
-                        secondary_index_building=len(search_settings) > 1,
-                        db_session=db_session,
-                    ):
-                        continue
-
-                    # using a task queue and only allowing one task per cc_pair/search_setting
-                    # prevents us from starving out certain attempts
-                    attempt_id = try_creating_indexing_task(
-                        self.app,
-                        cc_pair,
-                        search_settings_instance,
-                        False,
-                        db_session,
-                        r,
-                        tenant_id,
-                    )
-                    if attempt_id:
-                        task_logger.info(
-                            f"Indexing queued: cc_pair_id={cc_pair.id} index_attempt_id={attempt_id}"
-                        )
-                        tasks_created += 1
-    except SoftTimeLimitExceeded:
-        task_logger.info(
-            "Soft time limit exceeded, task is being terminated gracefully."
-        )
-    except Exception:
-        task_logger.exception("Unexpected exception")
-    finally:
-        if lock_beat.owned():
-            lock_beat.release()
-
-    return tasks_created
-
-
-def _should_index(
-    cc_pair: ConnectorCredentialPair,
-    last_index: IndexAttempt | None,
-    search_settings_instance: SearchSettings,
-    secondary_index_building: bool,
-    db_session: Session,
-) -> bool:
-    """Checks various global settings and past indexing attempts to determine if
-    we should try to start indexing the cc pair / search setting combination.
-
-    Note that tactical checks such as preventing overlap with a currently running task
-    are not handled here.
-
-    Return True if we should try to index, False if not.
-    """
-    connector = cc_pair.connector
-
-    # uncomment for debugging
-    # task_logger.info(f"_should_index: "
-    #                  f"cc_pair={cc_pair.id} "
-    #                  f"connector={cc_pair.connector_id} "
-    #                  f"refresh_freq={connector.refresh_freq}")
-
-    # don't kick off indexing for `NOT_APPLICABLE` sources
-    if connector.source == DocumentSource.NOT_APPLICABLE:
-        return False
-
-    # User can still manually create single indexing attempts via the UI for the
-    # currently in use index
-    if DISABLE_INDEX_UPDATE_ON_SWAP:
-        if (
-            search_settings_instance.status == IndexModelStatus.PRESENT
-            and secondary_index_building
-        ):
-            return False
-
-    # When switching over models, always index at least once
-    if search_settings_instance.status == IndexModelStatus.FUTURE:
-        if last_index:
-            # No new index if the last index attempt succeeded
-            # Once is enough. The model will never be able to swap otherwise.
-            if last_index.status == IndexingStatus.SUCCESS:
-                return False
-
-            # No new index if the last index attempt is waiting to start
-            if last_index.status == IndexingStatus.NOT_STARTED:
-                return False
-
-            # No new index if the last index attempt is running
-            if last_index.status == IndexingStatus.IN_PROGRESS:
-                return False
-        else:
-            if (
-                connector.id == 0 or connector.source == DocumentSource.INGESTION_API
-            ):  # Ingestion API
-                return False
-        return True
-
-    # If the connector is paused or is the ingestion API, don't index
-    # NOTE: during an embedding model switch over, the following logic
-    # is bypassed by the above check for a future model
-    if (
-        not cc_pair.status.is_active()
-        or connector.id == 0
-        or connector.source == DocumentSource.INGESTION_API
-    ):
-        return False
-
-    # if no attempt has ever occurred, we should index regardless of refresh_freq
-    if not last_index:
-        return True
-
-    if connector.refresh_freq is None:
-        return False
-
-    current_db_time = get_db_current_time(db_session)
-    time_since_index = current_db_time - last_index.time_updated
-    if time_since_index.total_seconds() < connector.refresh_freq:
-        return False
-
-    return True
-
-
-def try_creating_indexing_task(
-    celery_app: Celery,
-    cc_pair: ConnectorCredentialPair,
-    search_settings: SearchSettings,
-    reindex: bool,
-    db_session: Session,
-    r: Redis,
-    tenant_id: str | None,
-) -> int | None:
-    """Checks for any conditions that should block the indexing task from being
-    created, then creates the task.
-
-    Does not check for scheduling related conditions as this function
-    is used to trigger indexing immediately.
-    """
-
-    LOCK_TIMEOUT = 30
-
-    # we need to serialize any attempt to trigger indexing since it can be triggered
-    # either via celery beat or manually (API call)
-    lock = r.lock(
-        DANSWER_REDIS_FUNCTION_LOCK_PREFIX + "try_creating_indexing_task",
-        timeout=LOCK_TIMEOUT,
-    )
-
-    acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)
-    if not acquired:
-        return None
-
-    try:
-        rci = RedisConnectorIndexing(cc_pair.id, search_settings.id)
-
-        # skip if already indexing
-        if r.exists(rci.fence_key):
-            return None
-
-        # skip indexing if the cc_pair is deleting
-        rcd = RedisConnectorDeletion(cc_pair.id)
-        if r.exists(rcd.fence_key):
-            return None
-
-        db_session.refresh(cc_pair)
-        if cc_pair.status == ConnectorCredentialPairStatus.DELETING:
-            return None
-
-        # add a long running generator task to the queue
-        r.delete(rci.generator_complete_key)
-        r.delete(rci.taskset_key)
-
-        custom_task_id = f"{rci.generator_task_id_prefix}_{uuid4()}"
-
-        # set a basic fence to start
-        fence_value = RedisConnectorIndexingFenceData(
-            index_attempt_id=None,
-            started=None,
-            submitted=datetime.now(timezone.utc),
-            celery_task_id=None,
-        )
-        r.set(rci.fence_key, fence_value.model_dump_json())
-
-        # create the index attempt for tracking purposes
-        # code elsewhere checks for index attempts without an associated redis key
-        # and cleans them up
-        # therefore we must create the attempt and the task after the fence goes up
-        index_attempt_id = create_index_attempt(
-            cc_pair.id,
-            search_settings.id,
-            from_beginning=reindex,
-            db_session=db_session,
-        )
-
-        result = celery_app.send_task(
-            "connector_indexing_proxy_task",
-            kwargs=dict(
-                index_attempt_id=index_attempt_id,
-                cc_pair_id=cc_pair.id,
-                search_settings_id=search_settings.id,
-                tenant_id=tenant_id,
-            ),
-            queue=DanswerCeleryQueues.CONNECTOR_INDEXING,
-            task_id=custom_task_id,
-            priority=DanswerCeleryPriority.MEDIUM,
-        )
-        if not result:
-            raise RuntimeError("send_task for connector_indexing_proxy_task failed.")
-
-        # now fill out the fence with the rest of the data
-        fence_value.index_attempt_id = index_attempt_id
-        fence_value.celery_task_id = result.id
-        r.set(rci.fence_key, fence_value.model_dump_json())
-    except Exception:
-        r.delete(rci.fence_key)
-        task_logger.exception("Unexpected exception")
-        return None
-    finally:
-        if lock.owned():
-            lock.release()
-
-    return index_attempt_id
-
-
-@shared_task(name="connector_indexing_proxy_task", acks_late=False, track_started=True)
-def connector_indexing_proxy_task(
-    index_attempt_id: int,
-    cc_pair_id: int,
-    search_settings_id: int,
-    tenant_id: str | None,
-) -> None:
-    """celery tasks are forked, but forking is unstable.  This proxies work to a spawned task."""
-
-    client = SimpleJobClient()
-
-    job = client.submit(
-        connector_indexing_task,
-        index_attempt_id,
-        cc_pair_id,
-        search_settings_id,
-        tenant_id,
-        global_version.is_ee_version(),
-        pure=False,
-    )
-
-    if not job:
-        return
-
-    while True:
-        sleep(10)
-        with get_session_with_tenant(tenant_id) as db_session:
-            index_attempt = get_index_attempt(
-                db_session=db_session, index_attempt_id=index_attempt_id
-            )
-
-            # do nothing for ongoing jobs that haven't been stopped
-            if not job.done():
-                if not index_attempt:
-                    continue
-
-                if not index_attempt.is_finished():
-                    continue
-
-            if job.status == "error":
-                logger.error(job.exception())
-
-            job.release()
-            break
-
-    return
-
-
-def connector_indexing_task(
-    index_attempt_id: int,
-    cc_pair_id: int,
-    search_settings_id: int,
-    tenant_id: str | None,
-    is_ee: bool,
-) -> int | None:
-    """Indexing task. For a cc pair, this task pulls all document IDs from the source
-    and compares those IDs to locally stored documents and deletes all locally stored IDs missing
-    from the most recently pulled document ID list
-
-    acks_late must be set to False. Otherwise, celery's visibility timeout will
-    cause any task that runs longer than the timeout to be redispatched by the broker.
-    There appears to be no good workaround for this, so we need to handle redispatching
-    manually.
-
-    Returns None if the task did not run (possibly due to a conflict).
-    Otherwise, returns an int >= 0 representing the number of indexed docs.
-    """
-
-    attempt = None
-    n_final_progress = 0
-
-    r = get_redis_client(tenant_id=tenant_id)
-
-    rcd = RedisConnectorDeletion(cc_pair_id)
-    if r.exists(rcd.fence_key):
-        raise RuntimeError(
-            f"Indexing will not start because connector deletion is in progress: "
-            f"cc_pair={cc_pair_id} "
-            f"fence={rcd.fence_key}"
-        )
-
-    rcs = RedisConnectorStop(cc_pair_id)
-    if r.exists(rcs.fence_key):
-        raise RuntimeError(
-            f"Indexing will not start because a connector stop signal was detected: "
-            f"cc_pair={cc_pair_id} "
-            f"fence={rcs.fence_key}"
-        )
-
-    rci = RedisConnectorIndexing(cc_pair_id, search_settings_id)
-
-    while True:
-        # read related data and evaluate/print task progress
-        fence_value = cast(bytes, r.get(rci.fence_key))
-        if fence_value is None:
-            task_logger.info(
-                f"connector_indexing_task: fence_value not found: fence={rci.fence_key}"
-            )
-            raise RuntimeError(f"Fence not found: fence={rci.fence_key}")
-
-        try:
-            fence_json = fence_value.decode("utf-8")
-            fence_data = RedisConnectorIndexingFenceData.model_validate_json(
-                cast(str, fence_json)
-            )
-        except ValueError:
-            task_logger.exception(
-                f"connector_indexing_task: fence_data not decodeable: fence={rci.fence_key}"
-            )
-            raise
-
-        if fence_data.index_attempt_id is None or fence_data.celery_task_id is None:
-            task_logger.info(
-                f"connector_indexing_task - Waiting for fence: fence={rci.fence_key}"
-            )
-            sleep(1)
-            continue
-
-        task_logger.info(
-            f"connector_indexing_task - Fence found, continuing...: fence={rci.fence_key}"
-        )
-        break
-
-    lock = r.lock(
-        rci.generator_lock_key,
-        timeout=CELERY_INDEXING_LOCK_TIMEOUT,
-    )
-
-    acquired = lock.acquire(blocking=False)
-    if not acquired:
-        task_logger.warning(
-            f"Indexing task already running, exiting...: "
-            f"cc_pair={cc_pair_id} search_settings={search_settings_id}"
-        )
-        # r.set(rci.generator_complete_key, HTTPStatus.CONFLICT.value)
-        return None
-
-    fence_data.started = datetime.now(timezone.utc)
-    r.set(rci.fence_key, fence_data.model_dump_json())
-
-    try:
-        with get_session_with_tenant(tenant_id) as db_session:
-            attempt = get_index_attempt(db_session, index_attempt_id)
-            if not attempt:
-                raise ValueError(
-                    f"Index attempt not found: index_attempt={index_attempt_id}"
-                )
-
-            cc_pair = get_connector_credential_pair_from_id(
-                cc_pair_id=cc_pair_id,
-                db_session=db_session,
-            )
-
-            if not cc_pair:
-                raise ValueError(f"cc_pair not found: cc_pair={cc_pair_id}")
-
-            if not cc_pair.connector:
-                raise ValueError(
-                    f"Connector not found: cc_pair={cc_pair_id} connector={cc_pair.connector_id}"
-                )
-
-            if not cc_pair.credential:
-                raise ValueError(
-                    f"Credential not found: cc_pair={cc_pair_id} credential={cc_pair.credential_id}"
-                )
-
-            rci = RedisConnectorIndexing(cc_pair_id, search_settings_id)
-
-            # define a callback class
-            callback = RunIndexingCallback(
-                rcs.fence_key, rci.generator_progress_key, lock, r
-            )
-
-            run_indexing_entrypoint(
-                index_attempt_id,
-                tenant_id,
-                cc_pair_id,
-                is_ee,
-                callback=callback,
-            )
-
-            # get back the total number of indexed docs and return it
-            generator_progress_value = r.get(rci.generator_progress_key)
-            if generator_progress_value is not None:
-                try:
-                    n_final_progress = int(cast(int, generator_progress_value))
-                except ValueError:
-                    pass
-
-            r.set(rci.generator_complete_key, HTTPStatus.OK.value)
-    except Exception as e:
-        task_logger.exception(f"Indexing failed: cc_pair={cc_pair_id}")
-        if attempt:
-            with get_session_with_tenant(tenant_id) as db_session:
-                mark_attempt_failed(attempt, db_session, failure_reason=str(e))
-
-        r.delete(rci.generator_lock_key)
-        r.delete(rci.generator_progress_key)
-        r.delete(rci.taskset_key)
-        r.delete(rci.fence_key)
-        raise e
-    finally:
-        if lock.owned():
-            lock.release()
-
-    return n_final_progress

backend/danswer/background/celery/tasks/periodic/tasks.py

@@ -11,10 +11,10 @@ from sqlalchemy import inspect
 from sqlalchemy import text
 from sqlalchemy.orm import Session
 
-from danswer.background.celery.apps.app_base import task_logger
+from danswer.background.celery.celery_app import task_logger
 from danswer.configs.app_configs import JOB_TIMEOUT
 from danswer.configs.constants import PostgresAdvisoryLocks
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine  # type: ignore
 
 
 @shared_task(
@@ -23,7 +23,7 @@ from danswer.db.engine import get_session_with_tenant
     bind=True,
     base=AbortableTask,
 )
-def kombu_message_cleanup_task(self: Any, tenant_id: str | None) -> int:
+def kombu_message_cleanup_task(self: Any) -> int:
     """Runs periodically to clean up the kombu_message table"""
 
     # we will select messages older than this amount to clean up
@@ -35,7 +35,7 @@ def kombu_message_cleanup_task(self: Any, tenant_id: str | None) -> int:
     ctx["deleted"] = 0
     ctx["cleanup_age"] = KOMBU_MESSAGE_CLEANUP_AGE
     ctx["page_limit"] = KOMBU_MESSAGE_CLEANUP_PAGE_LIMIT
-    with get_session_with_tenant(tenant_id) as db_session:
+    with Session(get_sqlalchemy_engine()) as db_session:
         # Exit the task if we can't take the advisory lock
         result = db_session.execute(
             text("SELECT pg_try_advisory_lock(:id)"),

backend/danswer/background/celery/tasks/pruning/tasks.py

@@ -3,49 +3,39 @@ from datetime import timedelta
 from datetime import timezone
 from uuid import uuid4
 
-from celery import Celery
+import redis
 from celery import shared_task
-from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from redis import Redis
 from sqlalchemy.orm import Session
 
-from danswer.background.celery.apps.app_base import task_logger
-from danswer.background.celery.celery_redis import RedisConnectorDeletion
+from danswer.background.celery.celery_app import celery_app
+from danswer.background.celery.celery_app import task_logger
 from danswer.background.celery.celery_redis import RedisConnectorPruning
-from danswer.background.celery.celery_redis import RedisConnectorStop
 from danswer.background.celery.celery_utils import extract_ids_from_runnable_connector
-from danswer.background.celery.tasks.indexing.tasks import RunIndexingCallback
 from danswer.configs.app_configs import ALLOW_SIMULTANEOUS_PRUNING
 from danswer.configs.app_configs import JOB_TIMEOUT
-from danswer.configs.constants import CELERY_PRUNING_LOCK_TIMEOUT
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
-from danswer.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX
 from danswer.configs.constants import DanswerCeleryPriority
 from danswer.configs.constants import DanswerCeleryQueues
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.connectors.factory import instantiate_connector
 from danswer.connectors.models import InputType
 from danswer.db.connector_credential_pair import get_connector_credential_pair
-from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
 from danswer.db.connector_credential_pair import get_connector_credential_pairs
 from danswer.db.document import get_documents_for_connector_credential_pair
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.enums import ConnectorCredentialPairStatus
 from danswer.db.models import ConnectorCredentialPair
 from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
 
 
 @shared_task(
-    name="check_for_pruning",
+    name="check_for_prune_task_2",
     soft_time_limit=JOB_TIMEOUT,
-    bind=True,
 )
-def check_for_pruning(self: Task, *, tenant_id: str | None) -> None:
-    r = get_redis_client(tenant_id=tenant_id)
+def check_for_prune_task_2() -> None:
+    r = get_redis_client()
 
     lock_beat = r.lock(
         DanswerRedisLocks.CHECK_PRUNE_BEAT_LOCK,
@@ -57,29 +47,16 @@ def check_for_pruning(self: Task, *, tenant_id: str | None) -> None:
         if not lock_beat.acquire(blocking=False):
             return
 
-        cc_pair_ids: list[int] = []
-        with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             cc_pairs = get_connector_credential_pairs(db_session)
-            for cc_pair_entry in cc_pairs:
-                cc_pair_ids.append(cc_pair_entry.id)
-
-        for cc_pair_id in cc_pair_ids:
-            lock_beat.reacquire()
-            with get_session_with_tenant(tenant_id) as db_session:
-                cc_pair = get_connector_credential_pair_from_id(cc_pair_id, db_session)
-                if not cc_pair:
-                    continue
-
-                if not is_pruning_due(cc_pair, db_session, r):
-                    continue
-
-                tasks_created = try_creating_prune_generator_task(
-                    self.app, cc_pair, db_session, r, tenant_id
+            for cc_pair in cc_pairs:
+                tasks_created = ccpair_pruning_generator_task_creation_helper(
+                    cc_pair, db_session, r, lock_beat
                 )
                 if not tasks_created:
                     continue
 
-                task_logger.info(f"Pruning queued: cc_pair_id={cc_pair.id}")
+                task_logger.info(f"Pruning started: cc_pair_id={cc_pair.id}")
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
@@ -91,11 +68,12 @@ def check_for_pruning(self: Task, *, tenant_id: str | None) -> None:
             lock_beat.release()
 
 
-def is_pruning_due(
+def ccpair_pruning_generator_task_creation_helper(
     cc_pair: ConnectorCredentialPair,
     db_session: Session,
     r: Redis,
-) -> bool:
+    lock_beat: redis.lock.Lock,
+) -> int | None:
     """Returns an int if pruning is triggered.
     The int represents the number of prune tasks generated (in this case, only one
     because the task is a long running generator task.)
@@ -106,147 +84,84 @@ def is_pruning_due(
     try_creating_prune_generator_task.
     """
 
+    lock_beat.reacquire()
+
     # skip pruning if no prune frequency is set
     # pruning can still be forced via the API which will run a pruning task directly
     if not cc_pair.connector.prune_freq:
-        return False
-
-    # skip pruning if not active
-    if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:
-        return False
+        return None
 
     # skip pruning if the next scheduled prune time hasn't been reached yet
     last_pruned = cc_pair.last_pruned
     if not last_pruned:
-        if not cc_pair.last_successful_index_time:
-            # if we've never indexed, we can't prune
-            return False
-
-        # if never pruned, use the last time the connector indexed successfully
-        last_pruned = cc_pair.last_successful_index_time
+        # if never pruned, use the connector time created as the last_pruned time
+        last_pruned = cc_pair.connector.time_created
 
     next_prune = last_pruned + timedelta(seconds=cc_pair.connector.prune_freq)
     if datetime.now(timezone.utc) < next_prune:
-        return False
+        return None
 
-    return True
+    return try_creating_prune_generator_task(cc_pair, db_session, r)
 
 
 def try_creating_prune_generator_task(
-    celery_app: Celery,
     cc_pair: ConnectorCredentialPair,
     db_session: Session,
     r: Redis,
-    tenant_id: str | None,
 ) -> int | None:
     """Checks for any conditions that should block the pruning generator task from being
     created, then creates the task.
 
     Does not check for scheduling related conditions as this function
-    is used to trigger prunes immediately, e.g. via the web ui.
+    is used to trigger prunes immediately.
     """
 
     if not ALLOW_SIMULTANEOUS_PRUNING:
         for key in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
             return None
 
-    LOCK_TIMEOUT = 30
+    rcp = RedisConnectorPruning(cc_pair.id)
 
-    # we need to serialize starting pruning since it can be triggered either via
-    # celery beat or manually (API call)
-    lock = r.lock(
-        DANSWER_REDIS_FUNCTION_LOCK_PREFIX + "try_creating_prune_generator_task",
-        timeout=LOCK_TIMEOUT,
+    # skip pruning if already pruning
+    if r.exists(rcp.fence_key):
+        return None
+
+    # skip pruning if the cc_pair is deleting
+    db_session.refresh(cc_pair)
+    if cc_pair.status == ConnectorCredentialPairStatus.DELETING:
+        return None
+
+    # add a long running generator task to the queue
+    r.delete(rcp.generator_complete_key)
+    r.delete(rcp.taskset_key)
+
+    custom_task_id = f"{rcp.generator_task_id_prefix}_{uuid4()}"
+
+    celery_app.send_task(
+        "connector_pruning_generator_task",
+        kwargs=dict(
+            connector_id=cc_pair.connector_id, credential_id=cc_pair.credential_id
+        ),
+        queue=DanswerCeleryQueues.CONNECTOR_PRUNING,
+        task_id=custom_task_id,
+        priority=DanswerCeleryPriority.LOW,
     )
 
-    acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)
-    if not acquired:
-        return None
-
-    try:
-        rcp = RedisConnectorPruning(cc_pair.id)
-
-        # skip pruning if already pruning
-        if r.exists(rcp.fence_key):
-            return None
-
-        # skip pruning if the cc_pair is deleting
-        rcd = RedisConnectorDeletion(cc_pair.id)
-        if r.exists(rcd.fence_key):
-            return None
-
-        db_session.refresh(cc_pair)
-        if cc_pair.status == ConnectorCredentialPairStatus.DELETING:
-            return None
-
-        # add a long running generator task to the queue
-        r.delete(rcp.generator_complete_key)
-        r.delete(rcp.taskset_key)
-
-        custom_task_id = f"{rcp.generator_task_id_prefix}_{uuid4()}"
-
-        celery_app.send_task(
-            "connector_pruning_generator_task",
-            kwargs=dict(
-                cc_pair_id=cc_pair.id,
-                connector_id=cc_pair.connector_id,
-                credential_id=cc_pair.credential_id,
-                tenant_id=tenant_id,
-            ),
-            queue=DanswerCeleryQueues.CONNECTOR_PRUNING,
-            task_id=custom_task_id,
-            priority=DanswerCeleryPriority.LOW,
-        )
-
-        # set this only after all tasks have been added
-        r.set(rcp.fence_key, 1)
-    except Exception:
-        task_logger.exception("Unexpected exception")
-        return None
-    finally:
-        if lock.owned():
-            lock.release()
-
+    # set this only after all tasks have been added
+    r.set(rcp.fence_key, 1)
     return 1
 
 
-@shared_task(
-    name="connector_pruning_generator_task",
-    acks_late=False,
-    soft_time_limit=JOB_TIMEOUT,
-    track_started=True,
-    trail=False,
-    bind=True,
-)
-def connector_pruning_generator_task(
-    self: Task,
-    cc_pair_id: int,
-    connector_id: int,
-    credential_id: int,
-    tenant_id: str | None,
-) -> None:
+@shared_task(name="connector_pruning_generator_task", soft_time_limit=JOB_TIMEOUT)
+def connector_pruning_generator_task(connector_id: int, credential_id: int) -> None:
     """connector pruning task. For a cc pair, this task pulls all document IDs from the source
     and compares those IDs to locally stored documents and deletes all locally stored IDs missing
     from the most recently pulled document ID list"""
 
-    r = get_redis_client(tenant_id=tenant_id)
+    r = get_redis_client()
 
-    rcp = RedisConnectorPruning(cc_pair_id)
-
-    lock = r.lock(
-        DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{rcp._id}",
-        timeout=CELERY_PRUNING_LOCK_TIMEOUT,
-    )
-
-    acquired = lock.acquire(blocking=False)
-    if not acquired:
-        task_logger.warning(
-            f"Pruning task already running, exiting...: cc_pair={cc_pair_id}"
-        )
-        return None
-
-    try:
-        with get_session_with_tenant(tenant_id) as db_session:
+    with Session(get_sqlalchemy_engine()) as db_session:
+        try:
             cc_pair = get_connector_credential_pair(
                 db_session=db_session,
                 connector_id=connector_id,
@@ -255,26 +170,27 @@ def connector_pruning_generator_task(
 
             if not cc_pair:
                 task_logger.warning(
-                    f"cc_pair not found for {connector_id} {credential_id}"
+                    f"ccpair not found for {connector_id} {credential_id}"
                 )
                 return
 
+            rcp = RedisConnectorPruning(cc_pair.id)
+
+            # Define the callback function
+            def redis_increment_callback(amount: int) -> None:
+                r.incrby(rcp.generator_progress_key, amount)
+
             runnable_connector = instantiate_connector(
                 db_session,
                 cc_pair.connector.source,
-                InputType.SLIM_RETRIEVAL,
+                InputType.PRUNE,
                 cc_pair.connector.connector_specific_config,
                 cc_pair.credential,
             )
 
-            rcs = RedisConnectorStop(cc_pair_id)
-
-            callback = RunIndexingCallback(
-                rcs.fence_key, rcp.generator_progress_key, lock, r
-            )
             # a list of docs in the source
             all_connector_doc_ids: set[str] = extract_ids_from_runnable_connector(
-                runnable_connector, callback
+                runnable_connector, redis_increment_callback
             )
 
             # a list of docs in our local index
@@ -292,7 +208,7 @@ def connector_pruning_generator_task(
 
             task_logger.info(
                 f"Pruning set collected: "
-                f"cc_pair={cc_pair_id} "
+                f"cc_pair_id={cc_pair.id} "
                 f"docs_to_remove={len(doc_ids_to_remove)} "
                 f"doc_source={cc_pair.connector.source}"
             )
@@ -300,29 +216,24 @@ def connector_pruning_generator_task(
             rcp.documents_to_prune = set(doc_ids_to_remove)
 
             task_logger.info(
-                f"RedisConnectorPruning.generate_tasks starting. cc_pair_id={cc_pair_id}"
-            )
-            tasks_generated = rcp.generate_tasks(
-                self.app, db_session, r, None, tenant_id
+                f"RedisConnectorPruning.generate_tasks starting. cc_pair_id={cc_pair.id}"
             )
+            tasks_generated = rcp.generate_tasks(celery_app, db_session, r, None)
             if tasks_generated is None:
                 return None
 
             task_logger.info(
                 f"RedisConnectorPruning.generate_tasks finished. "
-                f"cc_pair={cc_pair_id} tasks_generated={tasks_generated}"
+                f"cc_pair_id={cc_pair.id} tasks_generated={tasks_generated}"
             )
 
             r.set(rcp.generator_complete_key, tasks_generated)
-    except Exception as e:
-        task_logger.exception(
-            f"Failed to run pruning: cc_pair={cc_pair_id} connector={connector_id}"
-        )
+        except Exception as e:
+            task_logger.exception(
+                f"Failed to run pruning for connector id {connector_id}."
+            )
 
-        r.delete(rcp.generator_progress_key)
-        r.delete(rcp.taskset_key)
-        r.delete(rcp.fence_key)
-        raise e
-    finally:
-        if lock.owned():
-            lock.release()
+            r.delete(rcp.generator_progress_key)
+            r.delete(rcp.taskset_key)
+            r.delete(rcp.fence_key)
+            raise e

backend/danswer/background/celery/tasks/shared/RedisConnectorDeletionFenceData.py

@@ -1,8 +0,0 @@
-from datetime import datetime
-
-from pydantic import BaseModel
-
-
-class RedisConnectorDeletionFenceData(BaseModel):
-    num_tasks: int | None
-    submitted: datetime

backend/danswer/background/celery/tasks/shared/RedisConnectorIndexingFenceData.py

@@ -1,10 +0,0 @@
-from datetime import datetime
-
-from pydantic import BaseModel
-
-
-class RedisConnectorIndexingFenceData(BaseModel):
-    index_attempt_id: int | None
-    started: datetime | None
-    submitted: datetime
-    celery_task_id: str | None

backend/danswer/background/celery/tasks/shared/tasks.py

@@ -1,38 +1,32 @@
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
+from sqlalchemy.orm import Session
 
 from danswer.access.access import get_access_for_document
-from danswer.background.celery.apps.app_base import task_logger
+from danswer.background.celery.celery_app import task_logger
 from danswer.db.document import delete_document_by_connector_credential_pair__no_commit
 from danswer.db.document import delete_documents_complete__no_commit
 from danswer.db.document import get_document
 from danswer.db.document import get_document_connector_count
-from danswer.db.document import mark_document_as_modified
 from danswer.db.document import mark_document_as_synced
 from danswer.db.document_set import fetch_document_sets_for_document
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.document_index.document_index_utils import get_both_index_names
 from danswer.document_index.factory import get_default_document_index
 from danswer.document_index.interfaces import VespaDocumentFields
 from danswer.server.documents.models import ConnectorCredentialPairIdentifier
 
-DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES = 3
-
 
 @shared_task(
     name="document_by_cc_pair_cleanup_task",
     bind=True,
     soft_time_limit=45,
     time_limit=60,
-    max_retries=DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES,
+    max_retries=3,
 )
 def document_by_cc_pair_cleanup_task(
-    self: Task,
-    document_id: str,
-    connector_id: int,
-    credential_id: int,
-    tenant_id: str | None,
+    self: Task, document_id: str, connector_id: int, credential_id: int
 ) -> bool:
     """A lightweight subtask used to clean up document to cc pair relationships.
     Created by connection deletion and connector pruning parent tasks."""
@@ -49,13 +43,10 @@ def document_by_cc_pair_cleanup_task(
     connector / credential pair from the access list
     (6) delete all relevant entries from postgres
     """
-    task_logger.info(f"tenant_id={tenant_id} document_id={document_id}")
+    task_logger.info(f"document_id={document_id}")
 
     try:
-        with get_session_with_tenant(tenant_id) as db_session:
-            action = "skip"
-            chunks_affected = 0
-
+        with Session(get_sqlalchemy_engine()) as db_session:
             curr_ind_name, sec_ind_name = get_both_index_names(db_session)
             document_index = get_default_document_index(
                 primary_index_name=curr_ind_name, secondary_index_name=sec_ind_name
@@ -65,16 +56,12 @@ def document_by_cc_pair_cleanup_task(
             if count == 1:
                 # count == 1 means this is the only remaining cc_pair reference to the doc
                 # delete it from vespa and the db
-                action = "delete"
-
-                chunks_affected = document_index.delete_single(document_id)
+                document_index.delete(doc_ids=[document_id])
                 delete_documents_complete__no_commit(
                     db_session=db_session,
                     document_ids=[document_id],
                 )
             elif count > 1:
-                action = "update"
-
                 # count > 1 means the document still has cc_pair references
                 doc = get_document(document_id, db_session)
                 if not doc:
@@ -97,9 +84,7 @@ def document_by_cc_pair_cleanup_task(
                 )
 
                 # update Vespa. OK if doc doesn't exist. Raises exception otherwise.
-                chunks_affected = document_index.update_single(
-                    document_id, fields=fields
-                )
+                document_index.update_single(document_id, fields=fields)
 
                 # there are still other cc_pair references to the doc, so just resync to Vespa
                 delete_document_by_connector_credential_pair__no_commit(
@@ -116,35 +101,13 @@ def document_by_cc_pair_cleanup_task(
                 pass
 
             db_session.commit()
-
-            task_logger.info(
-                f"tenant_id={tenant_id} "
-                f"document_id={document_id} "
-                f"action={action} "
-                f"refcount={count} "
-                f"chunks={chunks_affected}"
-            )
     except SoftTimeLimitExceeded:
-        task_logger.info(
-            f"SoftTimeLimitExceeded exception. tenant_id={tenant_id} doc_id={document_id}"
-        )
-        return False
+        task_logger.info(f"SoftTimeLimitExceeded exception. doc_id={document_id}")
     except Exception as e:
         task_logger.exception("Unexpected exception")
 
-        if self.request.retries < DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES:
-            # Still retrying. Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64
-            countdown = 2 ** (self.request.retries + 4)
-            self.retry(exc=e, countdown=countdown)
-        else:
-            # This is the last attempt! mark the document as dirty in the db so that it
-            # eventually gets fixed out of band via stale document reconciliation
-            task_logger.info(
-                f"Max retries reached. Marking doc as dirty for reconciliation: "
-                f"tenant_id={tenant_id} document_id={document_id}"
-            )
-            with get_session_with_tenant(tenant_id):
-                mark_document_as_modified(document_id, db_session)
-        return False
+        # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64
+        countdown = 2 ** (self.request.retries + 4)
+        self.retry(exc=e, countdown=countdown)
 
     return True

backend/danswer/background/celery/tasks/vespa/tasks.py

@@ -1,37 +1,23 @@
 import traceback
-from datetime import datetime
-from datetime import timezone
-from http import HTTPStatus
 from typing import cast
 
 import redis
-from celery import Celery
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
-from celery.result import AsyncResult
-from celery.states import READY_STATES
 from redis import Redis
 from sqlalchemy.orm import Session
 
 from danswer.access.access import get_access_for_document
-from danswer.background.celery.apps.app_base import task_logger
-from danswer.background.celery.celery_redis import celery_get_queue_length
+from danswer.background.celery.celery_app import celery_app
+from danswer.background.celery.celery_app import task_logger
 from danswer.background.celery.celery_redis import RedisConnectorCredentialPair
 from danswer.background.celery.celery_redis import RedisConnectorDeletion
-from danswer.background.celery.celery_redis import RedisConnectorIndexing
 from danswer.background.celery.celery_redis import RedisConnectorPruning
 from danswer.background.celery.celery_redis import RedisDocumentSet
 from danswer.background.celery.celery_redis import RedisUserGroup
-from danswer.background.celery.tasks.shared.RedisConnectorDeletionFenceData import (
-    RedisConnectorDeletionFenceData,
-)
-from danswer.background.celery.tasks.shared.RedisConnectorIndexingFenceData import (
-    RedisConnectorIndexingFenceData,
-)
 from danswer.configs.app_configs import JOB_TIMEOUT
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
-from danswer.configs.constants import DanswerCeleryQueues
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.db.connector import fetch_connector_by_id
 from danswer.db.connector import mark_ccpair_as_pruned
@@ -43,7 +29,6 @@ from danswer.db.connector_credential_pair import get_connector_credential_pair_f
 from danswer.db.connector_credential_pair import get_connector_credential_pairs
 from danswer.db.document import count_documents_by_needs_sync
 from danswer.db.document import get_document
-from danswer.db.document import get_document_ids_for_connector_credential_pair
 from danswer.db.document import mark_document_as_synced
 from danswer.db.document_set import delete_document_set
 from danswer.db.document_set import delete_document_set_cc_pair_relationship__no_commit
@@ -51,28 +36,20 @@ from danswer.db.document_set import fetch_document_sets
 from danswer.db.document_set import fetch_document_sets_for_document
 from danswer.db.document_set import get_document_set_by_id
 from danswer.db.document_set import mark_document_set_as_synced
-from danswer.db.engine import get_session_with_tenant
-from danswer.db.enums import IndexingStatus
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.index_attempt import delete_index_attempts
-from danswer.db.index_attempt import get_all_index_attempts_by_status
-from danswer.db.index_attempt import get_index_attempt
-from danswer.db.index_attempt import mark_attempt_failed
 from danswer.db.models import DocumentSet
-from danswer.db.models import IndexAttempt
+from danswer.db.models import UserGroup
 from danswer.document_index.document_index_utils import get_both_index_names
 from danswer.document_index.factory import get_default_document_index
-from danswer.document_index.interfaces import VespaDocumentFields
+from danswer.document_index.interfaces import UpdateRequest
 from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import setup_logger
 from danswer.utils.variable_functionality import fetch_versioned_implementation
 from danswer.utils.variable_functionality import (
     fetch_versioned_implementation_with_fallback,
 )
-from danswer.utils.variable_functionality import global_version
 from danswer.utils.variable_functionality import noop_fallback
 
-logger = setup_logger()
-
 
 # celery auto associates tasks created inside another task,
 # which bloats the result metadata considerably. trail=False prevents this.
@@ -80,13 +57,12 @@ logger = setup_logger()
     name="check_for_vespa_sync_task",
     soft_time_limit=JOB_TIMEOUT,
     trail=False,
-    bind=True,
 )
-def check_for_vespa_sync_task(self: Task, *, tenant_id: str | None) -> None:
+def check_for_vespa_sync_task() -> None:
     """Runs periodically to check if any document needs syncing.
     Generates sets of tasks for Celery if syncing is needed."""
 
-    r = get_redis_client(tenant_id=tenant_id)
+    r = get_redis_client()
 
     lock_beat = r.lock(
         DanswerRedisLocks.CHECK_VESPA_SYNC_BEAT_LOCK,
@@ -98,55 +74,34 @@ def check_for_vespa_sync_task(self: Task, *, tenant_id: str | None) -> None:
         if not lock_beat.acquire(blocking=False):
             return
 
-        with get_session_with_tenant(tenant_id) as db_session:
-            try_generate_stale_document_sync_tasks(
-                self.app, db_session, r, lock_beat, tenant_id
-            )
+        with Session(get_sqlalchemy_engine()) as db_session:
+            try_generate_stale_document_sync_tasks(db_session, r, lock_beat)
 
-        # region document set scan
-        document_set_ids: list[int] = []
-        with get_session_with_tenant(tenant_id) as db_session:
             # check if any document sets are not synced
             document_set_info = fetch_document_sets(
                 user_id=None, db_session=db_session, include_outdated=True
             )
-
             for document_set, _ in document_set_info:
-                document_set_ids.append(document_set.id)
-
-        for document_set_id in document_set_ids:
-            with get_session_with_tenant(tenant_id) as db_session:
                 try_generate_document_set_sync_tasks(
-                    self.app, document_set_id, db_session, r, lock_beat, tenant_id
+                    document_set, db_session, r, lock_beat
                 )
-        # endregion
 
-        # check if any user groups are not synced
-        if global_version.is_ee_version():
+            # check if any user groups are not synced
             try:
                 fetch_user_groups = fetch_versioned_implementation(
                     "danswer.db.user_group", "fetch_user_groups"
                 )
+
+                user_groups = fetch_user_groups(
+                    db_session=db_session, only_up_to_date=False
+                )
+                for usergroup in user_groups:
+                    try_generate_user_group_sync_tasks(
+                        usergroup, db_session, r, lock_beat
+                    )
             except ModuleNotFoundError:
                 # Always exceptions on the MIT version, which is expected
-                # We shouldn't actually get here if the ee version check works
                 pass
-            else:
-                usergroup_ids: list[int] = []
-                with get_session_with_tenant(tenant_id) as db_session:
-                    user_groups = fetch_user_groups(
-                        db_session=db_session, only_up_to_date=False
-                    )
-
-                    for usergroup in user_groups:
-                        usergroup_ids.append(usergroup.id)
-
-                for usergroup_id in usergroup_ids:
-                    with get_session_with_tenant(tenant_id) as db_session:
-                        try_generate_user_group_sync_tasks(
-                            self.app, usergroup_id, db_session, r, lock_beat, tenant_id
-                        )
-
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
@@ -159,11 +114,7 @@ def check_for_vespa_sync_task(self: Task, *, tenant_id: str | None) -> None:
 
 
 def try_generate_stale_document_sync_tasks(
-    celery_app: Celery,
-    db_session: Session,
-    r: Redis,
-    lock_beat: redis.lock.Lock,
-    tenant_id: str | None,
+    db_session: Session, r: Redis, lock_beat: redis.lock.Lock
 ) -> int | None:
     # the fence is up, do nothing
     if r.exists(RedisConnectorCredentialPair.get_fence_key()):
@@ -188,9 +139,7 @@ def try_generate_stale_document_sync_tasks(
     cc_pairs = get_connector_credential_pairs(db_session)
     for cc_pair in cc_pairs:
         rc = RedisConnectorCredentialPair(cc_pair.id)
-        tasks_generated = rc.generate_tasks(
-            celery_app, db_session, r, lock_beat, tenant_id
-        )
+        tasks_generated = rc.generate_tasks(celery_app, db_session, r, lock_beat)
 
         if tasks_generated is None:
             continue
@@ -214,16 +163,11 @@ def try_generate_stale_document_sync_tasks(
 
 
 def try_generate_document_set_sync_tasks(
-    celery_app: Celery,
-    document_set_id: int,
-    db_session: Session,
-    r: Redis,
-    lock_beat: redis.lock.Lock,
-    tenant_id: str | None,
+    document_set: DocumentSet, db_session: Session, r: Redis, lock_beat: redis.lock.Lock
 ) -> int | None:
     lock_beat.reacquire()
 
-    rds = RedisDocumentSet(document_set_id)
+    rds = RedisDocumentSet(document_set.id)
 
     # don't generate document set sync tasks if tasks are still pending
     if r.exists(rds.fence_key):
@@ -231,10 +175,7 @@ def try_generate_document_set_sync_tasks(
 
     # don't generate sync tasks if we're up to date
     # race condition with the monitor/cleanup function if we use a cached result!
-    document_set = get_document_set_by_id(db_session, document_set_id)
-    if not document_set:
-        return None
-
+    db_session.refresh(document_set)
     if document_set.is_up_to_date:
         return None
 
@@ -246,9 +187,7 @@ def try_generate_document_set_sync_tasks(
     )
 
     # Add all documents that need to be updated into the queue
-    tasks_generated = rds.generate_tasks(
-        celery_app, db_session, r, lock_beat, tenant_id
-    )
+    tasks_generated = rds.generate_tasks(celery_app, db_session, r, lock_beat)
     if tasks_generated is None:
         return None
 
@@ -269,30 +208,18 @@ def try_generate_document_set_sync_tasks(
 
 
 def try_generate_user_group_sync_tasks(
-    celery_app: Celery,
-    usergroup_id: int,
-    db_session: Session,
-    r: Redis,
-    lock_beat: redis.lock.Lock,
-    tenant_id: str | None,
+    usergroup: UserGroup, db_session: Session, r: Redis, lock_beat: redis.lock.Lock
 ) -> int | None:
     lock_beat.reacquire()
 
-    rug = RedisUserGroup(usergroup_id)
+    rug = RedisUserGroup(usergroup.id)
 
     # don't generate sync tasks if tasks are still pending
     if r.exists(rug.fence_key):
         return None
 
     # race condition with the monitor/cleanup function if we use a cached result!
-    fetch_user_group = fetch_versioned_implementation(
-        "danswer.db.user_group", "fetch_user_group"
-    )
-
-    usergroup = fetch_user_group(db_session, usergroup_id)
-    if not usergroup:
-        return None
-
+    db_session.refresh(usergroup)
     if usergroup.is_up_to_date:
         return None
 
@@ -303,9 +230,7 @@ def try_generate_user_group_sync_tasks(
     task_logger.info(
         f"RedisUserGroup.generate_tasks starting. usergroup_id={usergroup.id}"
     )
-    tasks_generated = rug.generate_tasks(
-        celery_app, db_session, r, lock_beat, tenant_id
-    )
+    tasks_generated = rug.generate_tasks(celery_app, db_session, r, lock_beat)
     if tasks_generated is None:
         return None
 
@@ -350,13 +275,11 @@ def monitor_document_set_taskset(
     key_bytes: bytes, r: Redis, db_session: Session
 ) -> None:
     fence_key = key_bytes.decode("utf-8")
-    document_set_id_str = RedisDocumentSet.get_id_from_fence_key(fence_key)
-    if document_set_id_str is None:
+    document_set_id = RedisDocumentSet.get_id_from_fence_key(fence_key)
+    if document_set_id is None:
         task_logger.warning(f"could not parse document set id from {fence_key}")
         return
 
-    document_set_id = int(document_set_id_str)
-
     rds = RedisDocumentSet(document_set_id)
 
     fence_value = r.get(rds.fence_key)
@@ -371,8 +294,7 @@ def monitor_document_set_taskset(
 
     count = cast(int, r.scard(rds.taskset_key))
     task_logger.info(
-        f"Document set sync progress: document_set={document_set_id} "
-        f"remaining={count} initial={initial_count}"
+        f"Document set sync progress: document_set_id={document_set_id} remaining={count} initial={initial_count}"
     )
     if count > 0:
         return
@@ -386,83 +308,58 @@ def monitor_document_set_taskset(
             # if there are no connectors, then delete the document set.
             delete_document_set(document_set_row=document_set, db_session=db_session)
             task_logger.info(
-                f"Successfully deleted document set: document_set={document_set_id}"
+                f"Successfully deleted document set with ID: '{document_set_id}'!"
             )
         else:
             mark_document_set_as_synced(document_set_id, db_session)
             task_logger.info(
-                f"Successfully synced document set: document_set={document_set_id}"
+                f"Successfully synced document set with ID: '{document_set_id}'!"
             )
 
     r.delete(rds.taskset_key)
     r.delete(rds.fence_key)
 
 
-def monitor_connector_deletion_taskset(
-    key_bytes: bytes, r: Redis, tenant_id: str | None
-) -> None:
+def monitor_connector_deletion_taskset(key_bytes: bytes, r: Redis) -> None:
     fence_key = key_bytes.decode("utf-8")
-    cc_pair_id_str = RedisConnectorDeletion.get_id_from_fence_key(fence_key)
-    if cc_pair_id_str is None:
+    cc_pair_id = RedisConnectorDeletion.get_id_from_fence_key(fence_key)
+    if cc_pair_id is None:
         task_logger.warning(f"could not parse cc_pair_id from {fence_key}")
         return
 
-    cc_pair_id = int(cc_pair_id_str)
-
     rcd = RedisConnectorDeletion(cc_pair_id)
 
-    # read related data and evaluate/print task progress
-    fence_value = cast(bytes, r.get(rcd.fence_key))
+    fence_value = r.get(rcd.fence_key)
     if fence_value is None:
         return
 
     try:
-        fence_json = fence_value.decode("utf-8")
-        fence_data = RedisConnectorDeletionFenceData.model_validate_json(
-            cast(str, fence_json)
-        )
+        initial_count = int(cast(int, fence_value))
     except ValueError:
-        task_logger.exception(
-            "monitor_ccpair_indexing_taskset: fence_data not decodeable."
-        )
-        raise
-
-    # the fence is setting up but isn't ready yet
-    if fence_data.num_tasks is None:
+        task_logger.error("The value is not an integer.")
         return
 
     count = cast(int, r.scard(rcd.taskset_key))
     task_logger.info(
-        f"Connector deletion progress: cc_pair={cc_pair_id} remaining={count} initial={fence_data.num_tasks}"
+        f"Connector deletion progress: cc_pair_id={cc_pair_id} remaining={count} initial={initial_count}"
     )
     if count > 0:
         return
 
-    with get_session_with_tenant(tenant_id) as db_session:
+    with Session(get_sqlalchemy_engine()) as db_session:
         cc_pair = get_connector_credential_pair_from_id(cc_pair_id, db_session)
         if not cc_pair:
             task_logger.warning(
-                f"Connector deletion - cc_pair not found: cc_pair={cc_pair_id}"
+                f"monitor_connector_deletion_taskset - cc_pair_id not found: cc_pair_id={cc_pair_id}"
             )
             return
 
         try:
-            doc_ids = get_document_ids_for_connector_credential_pair(
-                db_session, cc_pair.connector_id, cc_pair.credential_id
-            )
-            if len(doc_ids) > 0:
-                # if this happens, documents somehow got added while deletion was in progress. Likely a bug
-                # gating off pruning and indexing work before deletion starts
-                task_logger.warning(
-                    f"Connector deletion - documents still found after taskset completion: "
-                    f"cc_pair={cc_pair_id} num={len(doc_ids)}"
-                )
-
             # clean up the rest of the related Postgres entities
             # index attempts
             delete_index_attempts(
                 db_session=db_session,
-                cc_pair_id=cc_pair_id,
+                cc_pair_id=cc_pair.id,
             )
 
             # document sets
@@ -479,7 +376,7 @@ def monitor_connector_deletion_taskset(
                 noop_fallback,
             )
             cleanup_user_groups(
-                cc_pair_id=cc_pair_id,
+                cc_pair_id=cc_pair.id,
                 db_session=db_session,
             )
 
@@ -496,27 +393,26 @@ def monitor_connector_deletion_taskset(
             )
             if not connector or not len(connector.credentials):
                 task_logger.info(
-                    "Connector deletion - Found no credentials left for connector, deleting connector"
+                    "Found no credentials left for connector, deleting connector"
                 )
                 db_session.delete(connector)
             db_session.commit()
         except Exception as e:
-            db_session.rollback()
             stack_trace = traceback.format_exc()
             error_message = f"Error: {str(e)}\n\nStack Trace:\n{stack_trace}"
-            add_deletion_failure_message(db_session, cc_pair_id, error_message)
+            add_deletion_failure_message(db_session, cc_pair.id, error_message)
             task_logger.exception(
-                f"Connector deletion exceptioned: "
-                f"cc_pair={cc_pair_id} connector={cc_pair.connector_id} credential={cc_pair.credential_id}"
+                f"Failed to run connector_deletion. "
+                f"cc_pair_id={cc_pair_id} connector_id={cc_pair.connector_id} credential_id={cc_pair.credential_id}"
             )
             raise e
 
     task_logger.info(
-        f"Connector deletion succeeded: "
-        f"cc_pair={cc_pair_id} "
-        f"connector={cc_pair.connector_id} "
-        f"credential={cc_pair.credential_id} "
-        f"docs_deleted={fence_data.num_tasks}"
+        f"Successfully deleted cc_pair: "
+        f"cc_pair_id={cc_pair_id} "
+        f"connector_id={cc_pair.connector_id} "
+        f"credential_id={cc_pair.credential_id} "
+        f"docs_deleted={initial_count}"
     )
 
     r.delete(rcd.taskset_key)
@@ -527,15 +423,13 @@ def monitor_ccpair_pruning_taskset(
     key_bytes: bytes, r: Redis, db_session: Session
 ) -> None:
     fence_key = key_bytes.decode("utf-8")
-    cc_pair_id_str = RedisConnectorPruning.get_id_from_fence_key(fence_key)
-    if cc_pair_id_str is None:
+    cc_pair_id = RedisConnectorPruning.get_id_from_fence_key(fence_key)
+    if cc_pair_id is None:
         task_logger.warning(
-            f"monitor_ccpair_pruning_taskset: could not parse cc_pair_id from {fence_key}"
+            f"monitor_connector_pruning_taskset: could not parse cc_pair_id from {fence_key}"
         )
         return
 
-    cc_pair_id = int(cc_pair_id_str)
-
     rcp = RedisConnectorPruning(cc_pair_id)
 
     fence_value = r.get(rcp.fence_key)
@@ -559,7 +453,7 @@ def monitor_ccpair_pruning_taskset(
     if count > 0:
         return
 
-    mark_ccpair_as_pruned(int(cc_pair_id), db_session)
+    mark_ccpair_as_pruned(cc_pair_id, db_session)
     task_logger.info(
         f"Successfully pruned connector credential pair. cc_pair_id={cc_pair_id}"
     )
@@ -570,136 +464,18 @@ def monitor_ccpair_pruning_taskset(
     r.delete(rcp.fence_key)
 
 
-def monitor_ccpair_indexing_taskset(
-    key_bytes: bytes, r: Redis, db_session: Session
-) -> None:
-    # if the fence doesn't exist, there's nothing to do
-    fence_key = key_bytes.decode("utf-8")
-    composite_id = RedisConnectorIndexing.get_id_from_fence_key(fence_key)
-    if composite_id is None:
-        task_logger.warning(
-            f"monitor_ccpair_indexing_taskset: could not parse composite_id from {fence_key}"
-        )
-        return
-
-    # parse out metadata and initialize the helper class with it
-    parts = composite_id.split("/")
-    if len(parts) != 2:
-        return
-
-    cc_pair_id = int(parts[0])
-    search_settings_id = int(parts[1])
-
-    rci = RedisConnectorIndexing(cc_pair_id, search_settings_id)
-
-    # read related data and evaluate/print task progress
-    fence_value = cast(bytes, r.get(rci.fence_key))
-    if fence_value is None:
-        return
-
-    try:
-        fence_json = fence_value.decode("utf-8")
-        fence_data = RedisConnectorIndexingFenceData.model_validate_json(
-            cast(str, fence_json)
-        )
-    except ValueError:
-        task_logger.exception(
-            "monitor_ccpair_indexing_taskset: fence_data not decodeable."
-        )
-        raise
-
-    elapsed_submitted = datetime.now(timezone.utc) - fence_data.submitted
-
-    generator_progress_value = r.get(rci.generator_progress_key)
-    if generator_progress_value is not None:
-        try:
-            progress_count = int(cast(int, generator_progress_value))
-
-            task_logger.info(
-                f"Connector indexing progress: cc_pair_id={cc_pair_id} "
-                f"search_settings_id={search_settings_id} "
-                f"progress={progress_count} "
-                f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
-            )
-        except ValueError:
-            task_logger.error(
-                "monitor_ccpair_indexing_taskset: generator_progress_value is not an integer."
-            )
-
-    if fence_data.index_attempt_id is None or fence_data.celery_task_id is None:
-        # the task is still setting up
-        return
-
-    # Read result state BEFORE generator_complete_key to avoid a race condition
-    # never use any blocking methods on the result from inside a task!
-    result: AsyncResult = AsyncResult(fence_data.celery_task_id)
-    result_state = result.state
-
-    generator_complete_value = r.get(rci.generator_complete_key)
-    if generator_complete_value is None:
-        if result_state in READY_STATES:
-            # IF the task state is READY, THEN generator_complete should be set
-            # if it isn't, then the worker crashed
-            task_logger.info(
-                f"Connector indexing aborted: "
-                f"cc_pair_id={cc_pair_id} "
-                f"search_settings_id={search_settings_id} "
-                f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
-            )
-
-            index_attempt = get_index_attempt(db_session, fence_data.index_attempt_id)
-            if index_attempt:
-                mark_attempt_failed(
-                    index_attempt=index_attempt,
-                    db_session=db_session,
-                    failure_reason="Connector indexing aborted or exceptioned.",
-                )
-
-            r.delete(rci.generator_lock_key)
-            r.delete(rci.taskset_key)
-            r.delete(rci.generator_progress_key)
-            r.delete(rci.generator_complete_key)
-            r.delete(rci.fence_key)
-        return
-
-    status_enum = HTTPStatus.INTERNAL_SERVER_ERROR
-    try:
-        status_value = int(cast(int, generator_complete_value))
-        status_enum = HTTPStatus(status_value)
-    except ValueError:
-        task_logger.error(
-            f"monitor_ccpair_indexing_taskset: "
-            f"generator_complete_value=f{generator_complete_value} could not be parsed."
-        )
-
-    task_logger.info(
-        f"Connector indexing finished: cc_pair_id={cc_pair_id} "
-        f"search_settings_id={search_settings_id} "
-        f"status={status_enum.name} "
-        f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
-    )
-
-    r.delete(rci.generator_lock_key)
-    r.delete(rci.taskset_key)
-    r.delete(rci.generator_progress_key)
-    r.delete(rci.generator_complete_key)
-    r.delete(rci.fence_key)
-
-
-@shared_task(name="monitor_vespa_sync", soft_time_limit=300, bind=True)
-def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
+@shared_task(name="monitor_vespa_sync", soft_time_limit=300)
+def monitor_vespa_sync() -> None:
     """This is a celery beat task that monitors and finalizes metadata sync tasksets.
     It scans for fence values and then gets the counts of any associated tasksets.
     If the count is 0, that means all tasks finished and we should clean up.
 
     This task lock timeout is CELERY_METADATA_SYNC_BEAT_LOCK_TIMEOUT seconds, so don't
     do anything too expensive in this function!
-
-    Returns True if the task actually did work, False
     """
-    r = get_redis_client(tenant_id=tenant_id)
+    r = get_redis_client()
 
-    lock_beat: redis.lock.Lock = r.lock(
+    lock_beat = r.lock(
         DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK,
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
@@ -707,90 +483,31 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     try:
         # prevent overlapping tasks
         if not lock_beat.acquire(blocking=False):
-            return False
+            return
 
-        # print current queue lengths
-        r_celery = self.app.broker_connection().channel().client  # type: ignore
-        n_celery = celery_get_queue_length("celery", r)
-        n_indexing = celery_get_queue_length(
-            DanswerCeleryQueues.CONNECTOR_INDEXING, r_celery
-        )
-        n_sync = celery_get_queue_length(
-            DanswerCeleryQueues.VESPA_METADATA_SYNC, r_celery
-        )
-        n_deletion = celery_get_queue_length(
-            DanswerCeleryQueues.CONNECTOR_DELETION, r_celery
-        )
-        n_pruning = celery_get_queue_length(
-            DanswerCeleryQueues.CONNECTOR_PRUNING, r_celery
-        )
-
-        task_logger.info(
-            f"Queue lengths: celery={n_celery} "
-            f"indexing={n_indexing} "
-            f"sync={n_sync} "
-            f"deletion={n_deletion} "
-            f"pruning={n_pruning}"
-        )
-
-        # do some cleanup before clearing fences
-        # check the db for any outstanding index attempts
-        with get_session_with_tenant(tenant_id) as db_session:
-            attempts: list[IndexAttempt] = []
-            attempts.extend(
-                get_all_index_attempts_by_status(IndexingStatus.NOT_STARTED, db_session)
-            )
-            attempts.extend(
-                get_all_index_attempts_by_status(IndexingStatus.IN_PROGRESS, db_session)
-            )
-
-            for a in attempts:
-                # if attempts exist in the db but we don't detect them in redis, mark them as failed
-                rci = RedisConnectorIndexing(
-                    a.connector_credential_pair_id, a.search_settings_id
-                )
-                failure_reason = f"Unknown index attempt {a.id}. Might be left over from a process restart."
-                if not r.exists(rci.fence_key):
-                    mark_attempt_failed(a, db_session, failure_reason=failure_reason)
-
-        lock_beat.reacquire()
         if r.exists(RedisConnectorCredentialPair.get_fence_key()):
             monitor_connector_taskset(r)
 
-        lock_beat.reacquire()
         for key_bytes in r.scan_iter(RedisConnectorDeletion.FENCE_PREFIX + "*"):
-            lock_beat.reacquire()
-            monitor_connector_deletion_taskset(key_bytes, r, tenant_id)
+            monitor_connector_deletion_taskset(key_bytes, r)
 
-        lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
-            lock_beat.reacquire()
-            with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
+            for key_bytes in r.scan_iter(RedisDocumentSet.FENCE_PREFIX + "*"):
                 monitor_document_set_taskset(key_bytes, r, db_session)
 
-        lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
-            lock_beat.reacquire()
-            monitor_usergroup_taskset = fetch_versioned_implementation_with_fallback(
-                "danswer.background.celery.tasks.vespa.tasks",
-                "monitor_usergroup_taskset",
-                noop_fallback,
-            )
-            with get_session_with_tenant(tenant_id) as db_session:
+            for key_bytes in r.scan_iter(RedisUserGroup.FENCE_PREFIX + "*"):
+                monitor_usergroup_taskset = (
+                    fetch_versioned_implementation_with_fallback(
+                        "danswer.background.celery.tasks.vespa.tasks",
+                        "monitor_usergroup_taskset",
+                        noop_fallback,
+                    )
+                )
                 monitor_usergroup_taskset(key_bytes, r, db_session)
 
-        lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
-            lock_beat.reacquire()
-            with get_session_with_tenant(tenant_id) as db_session:
+            for key_bytes in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
                 monitor_ccpair_pruning_taskset(key_bytes, r, db_session)
 
-        lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorIndexing.FENCE_PREFIX + "*"):
-            lock_beat.reacquire()
-            with get_session_with_tenant(tenant_id) as db_session:
-                monitor_ccpair_indexing_taskset(key_bytes, r, db_session)
-
         # uncomment for debugging if needed
         # r_celery = celery_app.broker_connection().channel().client
         # length = celery_get_queue_length(DanswerCeleryQueues.VESPA_METADATA_SYNC, r_celery)
@@ -803,8 +520,6 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
         if lock_beat.owned():
             lock_beat.release()
 
-    return True
-
 
 @shared_task(
     name="vespa_metadata_sync_task",
@@ -813,13 +528,11 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     time_limit=60,
     max_retries=3,
 )
-def vespa_metadata_sync_task(
-    self: Task, document_id: str, tenant_id: str | None
-) -> bool:
+def vespa_metadata_sync_task(self: Task, document_id: str) -> bool:
     task_logger.info(f"document_id={document_id}")
 
     try:
-        with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             curr_ind_name, sec_ind_name = get_both_index_names(db_session)
             document_index = get_default_document_index(
                 primary_index_name=curr_ind_name, secondary_index_name=sec_ind_name
@@ -837,24 +550,20 @@ def vespa_metadata_sync_task(
             doc_access = get_access_for_document(
                 document_id=document_id, db_session=db_session
             )
-
-            fields = VespaDocumentFields(
+            update_request = UpdateRequest(
+                document_ids=[document_id],
                 document_sets=update_doc_sets,
                 access=doc_access,
                 boost=doc.boost,
                 hidden=doc.hidden,
             )
 
-            # update Vespa. OK if doc doesn't exist. Raises exception otherwise.
-            chunks_affected = document_index.update_single(document_id, fields=fields)
+            # update Vespa
+            document_index.update(update_requests=[update_request])
 
             # update db last. Worst case = we crash right before this and
             # the sync might repeat again later
             mark_document_as_synced(document_id, db_session)
-
-            task_logger.info(
-                f"document_id={document_id} action=sync chunks={chunks_affected}"
-            )
     except SoftTimeLimitExceeded:
         task_logger.info(f"SoftTimeLimitExceeded exception. doc_id={document_id}")
     except Exception as e:

backend/danswer/background/celery/versioned_apps/beat.py

@@ -1,8 +0,0 @@
-"""Factory stub for running celery worker / celery beat."""
-from danswer.utils.variable_functionality import fetch_versioned_implementation
-from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
-
-set_is_ee_based_on_env_variable()
-app = fetch_versioned_implementation(
-    "danswer.background.celery.apps.beat", "celery_app"
-)

backend/danswer/background/celery/versioned_apps/heavy.py

@@ -1,17 +0,0 @@
-"""Factory stub for running celery worker / celery beat.
-This code is different from the primary/beat stubs because there is no EE version to
-fetch. Port over the code in those files if we add an EE version of this worker."""
-from celery import Celery
-
-from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
-
-set_is_ee_based_on_env_variable()
-
-
-def get_app() -> Celery:
-    from danswer.background.celery.apps.heavy import celery_app
-
-    return celery_app
-
-
-app = get_app()

backend/danswer/background/celery/versioned_apps/indexing.py

@@ -1,17 +0,0 @@
-"""Factory stub for running celery worker / celery beat.
-This code is different from the primary/beat stubs because there is no EE version to
-fetch. Port over the code in those files if we add an EE version of this worker."""
-from celery import Celery
-
-from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
-
-set_is_ee_based_on_env_variable()
-
-
-def get_app() -> Celery:
-    from danswer.background.celery.apps.indexing import celery_app
-
-    return celery_app
-
-
-app = get_app()

backend/danswer/background/celery/versioned_apps/light.py

@@ -1,17 +0,0 @@
-"""Factory stub for running celery worker / celery beat.
-This code is different from the primary/beat stubs because there is no EE version to
-fetch. Port over the code in those files if we add an EE version of this worker."""
-from celery import Celery
-
-from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
-
-set_is_ee_based_on_env_variable()
-
-
-def get_app() -> Celery:
-    from danswer.background.celery.apps.light import celery_app
-
-    return celery_app
-
-
-app = get_app()

backend/danswer/background/indexing/run_indexing.py

@@ -1,7 +1,5 @@
 import time
 import traceback
-from abc import ABC
-from abc import abstractmethod
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
@@ -19,12 +17,13 @@ from danswer.connectors.models import IndexAttemptMetadata
 from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
 from danswer.db.connector_credential_pair import get_last_successful_attempt_time
 from danswer.db.connector_credential_pair import update_connector_credential_pair
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.enums import ConnectorCredentialPairStatus
+from danswer.db.index_attempt import get_index_attempt
 from danswer.db.index_attempt import mark_attempt_failed
+from danswer.db.index_attempt import mark_attempt_in_progress
 from danswer.db.index_attempt import mark_attempt_partially_succeeded
 from danswer.db.index_attempt import mark_attempt_succeeded
-from danswer.db.index_attempt import transition_attempt_to_in_progress
 from danswer.db.index_attempt import update_docs_indexed
 from danswer.db.models import IndexAttempt
 from danswer.db.models import IndexingStatus
@@ -42,25 +41,11 @@ logger = setup_logger()
 INDEXING_TRACER_NUM_PRINT_ENTRIES = 5
 
 
-class RunIndexingCallbackInterface(ABC):
-    """Defines a callback interface to be passed to
-    to run_indexing_entrypoint."""
-
-    @abstractmethod
-    def should_stop(self) -> bool:
-        """Signal to stop the looping function in flight."""
-
-    @abstractmethod
-    def progress(self, amount: int) -> None:
-        """Send progress updates to the caller."""
-
-
 def _get_connector_runner(
     db_session: Session,
     attempt: IndexAttempt,
     start_time: datetime,
     end_time: datetime,
-    tenant_id: str | None,
 ) -> ConnectorRunner:
     """
     NOTE: `start_time` and `end_time` are only used for poll connectors
@@ -78,7 +63,6 @@ def _get_connector_runner(
             input_type=task,
             connector_specific_config=attempt.connector_credential_pair.connector.connector_specific_config,
             credential=attempt.connector_credential_pair.credential,
-            tenant_id=tenant_id,
         )
     except Exception as e:
         logger.exception(f"Unable to instantiate connector due to {e}")
@@ -105,16 +89,11 @@ def _get_connector_runner(
 def _run_indexing(
     db_session: Session,
     index_attempt: IndexAttempt,
-    tenant_id: str | None,
-    callback: RunIndexingCallbackInterface | None = None,
 ) -> None:
     """
     1. Get documents which are either new or updated from specified application
     2. Embed and index these documents into the chosen datastore (vespa)
     3. Updates Postgres to record the indexed documents + the outcome of this run
-
-    TODO: do not change index attempt statuses here ... instead, set signals in redis
-    and allow the monitor function to clean them up
     """
     start_time = time.time()
 
@@ -150,7 +129,6 @@ def _run_indexing(
             or (search_settings.status == IndexModelStatus.FUTURE)
         ),
         db_session=db_session,
-        tenant_id=tenant_id,
     )
 
     db_cc_pair = index_attempt.connector_credential_pair
@@ -207,7 +185,6 @@ def _run_indexing(
                 attempt=index_attempt,
                 start_time=window_start,
                 end_time=window_end,
-                tenant_id=tenant_id,
             )
 
             all_connector_doc_ids: set[str] = set()
@@ -220,12 +197,7 @@ def _run_indexing(
                 # index being built. We want to populate it even for paused connectors
                 # Often paused connectors are sources that aren't updated frequently but the
                 # contents still need to be initially pulled.
-                if callback:
-                    if callback.should_stop():
-                        raise RuntimeError("Connector stop signal detected")
-
-                # TODO: should we move this into the above callback instead?
-                db_session.refresh(db_cc_pair)
+                db_session.refresh(db_connector)
                 if (
                     (
                         db_cc_pair.status == ConnectorCredentialPairStatus.PAUSED
@@ -240,9 +212,7 @@ def _run_indexing(
                 db_session.refresh(index_attempt)
                 if index_attempt.status != IndexingStatus.IN_PROGRESS:
                     # Likely due to user manually disabling it or model swap
-                    raise RuntimeError(
-                        f"Index Attempt was canceled, status is {index_attempt.status}"
-                    )
+                    raise RuntimeError("Index Attempt was canceled")
 
                 batch_description = []
                 for doc in doc_batch:
@@ -262,8 +232,6 @@ def _run_indexing(
                 logger.debug(f"Indexing batch of documents: {batch_description}")
 
                 index_attempt_md.batch_num = batch_num + 1  # use 1-index for this
-
-                # real work happens here!
                 new_docs, total_batch_chunks = indexing_pipeline(
                     document_batch=doc_batch,
                     index_attempt_metadata=index_attempt_md,
@@ -282,9 +250,6 @@ def _run_indexing(
                 # be inaccurate
                 db_session.commit()
 
-                if callback:
-                    callback.progress(len(doc_batch))
-
                 # This new value is updated every batch, so UI can refresh per batch update
                 update_docs_indexed(
                     db_session=db_session,
@@ -408,13 +373,40 @@ def _run_indexing(
         )
 
 
+def _prepare_index_attempt(db_session: Session, index_attempt_id: int) -> IndexAttempt:
+    # make sure that the index attempt can't change in between checking the
+    # status and marking it as in_progress. This setting will be discarded
+    # after the next commit:
+    # https://docs.sqlalchemy.org/en/20/orm/session_transaction.html#setting-isolation-for-individual-transactions
+    db_session.connection(execution_options={"isolation_level": "SERIALIZABLE"})  # type: ignore
+
+    attempt = get_index_attempt(
+        db_session=db_session,
+        index_attempt_id=index_attempt_id,
+    )
+
+    if attempt is None:
+        raise RuntimeError(f"Unable to find IndexAttempt for ID '{index_attempt_id}'")
+
+    if attempt.status != IndexingStatus.NOT_STARTED:
+        raise RuntimeError(
+            f"Indexing attempt with ID '{index_attempt_id}' is not in NOT_STARTED status. "
+            f"Current status is '{attempt.status}'."
+        )
+
+    # only commit once, to make sure this all happens in a single transaction
+    mark_attempt_in_progress(attempt, db_session)
+
+    return attempt
+
+
 def run_indexing_entrypoint(
-    index_attempt_id: int,
-    tenant_id: str | None,
-    connector_credential_pair_id: int,
-    is_ee: bool = False,
-    callback: RunIndexingCallbackInterface | None = None,
+    index_attempt_id: int, connector_credential_pair_id: int, is_ee: bool = False
 ) -> None:
+    """Entrypoint for indexing run when using dask distributed.
+    Wraps the actual logic in a `try` block so that we can catch any exceptions
+    and mark the attempt as failed."""
+
     try:
         if is_ee:
             global_version.set_ee()
@@ -424,29 +416,26 @@ def run_indexing_entrypoint(
         IndexAttemptSingleton.set_cc_and_index_id(
             index_attempt_id, connector_credential_pair_id
         )
-        with get_session_with_tenant(tenant_id) as db_session:
-            attempt = transition_attempt_to_in_progress(index_attempt_id, db_session)
+
+        with Session(get_sqlalchemy_engine()) as db_session:
+            # make sure that it is valid to run this indexing attempt + mark it
+            # as in progress
+            attempt = _prepare_index_attempt(db_session, index_attempt_id)
 
             logger.info(
-                f"Indexing starting for tenant {tenant_id}: "
-                if tenant_id is not None
-                else ""
-                + f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"Indexing starting: "
+                f"connector='{attempt.connector_credential_pair.connector.name}' "
                 f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )
 
-            _run_indexing(db_session, attempt, tenant_id, callback)
+            _run_indexing(db_session, attempt)
 
             logger.info(
-                f"Indexing finished for tenant {tenant_id}: "
-                if tenant_id is not None
-                else ""
-                + f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"Indexing finished: "
+                f"connector='{attempt.connector_credential_pair.connector.name}' "
                 f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )
     except Exception as e:
-        logger.exception(
-            f"Indexing job with ID '{index_attempt_id}' for tenant {tenant_id} failed due to {e}"
-        )
+        logger.exception(f"Indexing job with ID '{index_attempt_id}' failed due to {e}")

backend/danswer/background/update.py

@@ -0,0 +1,495 @@
+import logging
+import time
+from datetime import datetime
+
+import dask
+from dask.distributed import Client
+from dask.distributed import Future
+from distributed import LocalCluster
+from sqlalchemy.orm import Session
+
+from danswer.background.indexing.dask_utils import ResourceLogger
+from danswer.background.indexing.job_client import SimpleJob
+from danswer.background.indexing.job_client import SimpleJobClient
+from danswer.background.indexing.run_indexing import run_indexing_entrypoint
+from danswer.configs.app_configs import CLEANUP_INDEXING_JOBS_TIMEOUT
+from danswer.configs.app_configs import DASK_JOB_CLIENT_ENABLED
+from danswer.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP
+from danswer.configs.app_configs import NUM_INDEXING_WORKERS
+from danswer.configs.app_configs import NUM_SECONDARY_INDEXING_WORKERS
+from danswer.configs.constants import DocumentSource
+from danswer.configs.constants import POSTGRES_INDEXER_APP_NAME
+from danswer.db.connector import fetch_connectors
+from danswer.db.connector_credential_pair import fetch_connector_credential_pairs
+from danswer.db.engine import get_db_current_time
+from danswer.db.engine import get_sqlalchemy_engine
+from danswer.db.engine import SqlEngine
+from danswer.db.index_attempt import create_index_attempt
+from danswer.db.index_attempt import get_index_attempt
+from danswer.db.index_attempt import get_inprogress_index_attempts
+from danswer.db.index_attempt import get_last_attempt_for_cc_pair
+from danswer.db.index_attempt import get_not_started_index_attempts
+from danswer.db.index_attempt import mark_attempt_failed
+from danswer.db.models import ConnectorCredentialPair
+from danswer.db.models import IndexAttempt
+from danswer.db.models import IndexingStatus
+from danswer.db.models import IndexModelStatus
+from danswer.db.models import SearchSettings
+from danswer.db.search_settings import get_current_search_settings
+from danswer.db.search_settings import get_secondary_search_settings
+from danswer.db.swap_index import check_index_swap
+from danswer.natural_language_processing.search_nlp_models import EmbeddingModel
+from danswer.natural_language_processing.search_nlp_models import warm_up_bi_encoder
+from danswer.utils.logger import setup_logger
+from danswer.utils.variable_functionality import global_version
+from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
+from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
+from shared_configs.configs import LOG_LEVEL
+from shared_configs.configs import MODEL_SERVER_PORT
+
+logger = setup_logger()
+
+# If the indexing dies, it's most likely due to resource constraints,
+# restarting just delays the eventual failure, not useful to the user
+dask.config.set({"distributed.scheduler.allowed-failures": 0})
+
+_UNEXPECTED_STATE_FAILURE_REASON = (
+    "Stopped mid run, likely due to the background process being killed"
+)
+
+
+def _should_create_new_indexing(
+    cc_pair: ConnectorCredentialPair,
+    last_index: IndexAttempt | None,
+    search_settings_instance: SearchSettings,
+    secondary_index_building: bool,
+    db_session: Session,
+) -> bool:
+    connector = cc_pair.connector
+
+    # don't kick off indexing for `NOT_APPLICABLE` sources
+    if connector.source == DocumentSource.NOT_APPLICABLE:
+        return False
+
+    # User can still manually create single indexing attempts via the UI for the
+    # currently in use index
+    if DISABLE_INDEX_UPDATE_ON_SWAP:
+        if (
+            search_settings_instance.status == IndexModelStatus.PRESENT
+            and secondary_index_building
+        ):
+            return False
+
+    # When switching over models, always index at least once
+    if search_settings_instance.status == IndexModelStatus.FUTURE:
+        if last_index:
+            # No new index if the last index attempt succeeded
+            # Once is enough. The model will never be able to swap otherwise.
+            if last_index.status == IndexingStatus.SUCCESS:
+                return False
+
+            # No new index if the last index attempt is waiting to start
+            if last_index.status == IndexingStatus.NOT_STARTED:
+                return False
+
+            # No new index if the last index attempt is running
+            if last_index.status == IndexingStatus.IN_PROGRESS:
+                return False
+        else:
+            if (
+                connector.id == 0 or connector.source == DocumentSource.INGESTION_API
+            ):  # Ingestion API
+                return False
+        return True
+
+    # If the connector is paused or is the ingestion API, don't index
+    # NOTE: during an embedding model switch over, the following logic
+    # is bypassed by the above check for a future model
+    if (
+        not cc_pair.status.is_active()
+        or connector.id == 0
+        or connector.source == DocumentSource.INGESTION_API
+    ):
+        return False
+
+    if not last_index:
+        return True
+
+    if connector.refresh_freq is None:
+        return False
+
+    # Only one scheduled/ongoing job per connector at a time
+    # this prevents cases where
+    # (1) the "latest" index_attempt is scheduled so we show
+    #     that in the UI despite another index_attempt being in-progress
+    # (2) multiple scheduled index_attempts at a time
+    if (
+        last_index.status == IndexingStatus.NOT_STARTED
+        or last_index.status == IndexingStatus.IN_PROGRESS
+    ):
+        return False
+
+    current_db_time = get_db_current_time(db_session)
+    time_since_index = current_db_time - last_index.time_updated
+    return time_since_index.total_seconds() >= connector.refresh_freq
+
+
+def _mark_run_failed(
+    db_session: Session, index_attempt: IndexAttempt, failure_reason: str
+) -> None:
+    """Marks the `index_attempt` row as failed + updates the `
+    connector_credential_pair` to reflect that the run failed"""
+    logger.warning(
+        f"Marking in-progress attempt 'connector: {index_attempt.connector_credential_pair.connector_id}, "
+        f"credential: {index_attempt.connector_credential_pair.credential_id}' as failed due to {failure_reason}"
+    )
+    mark_attempt_failed(
+        index_attempt=index_attempt,
+        db_session=db_session,
+        failure_reason=failure_reason,
+    )
+
+
+"""Main funcs"""
+
+
+def create_indexing_jobs(existing_jobs: dict[int, Future | SimpleJob]) -> None:
+    """Creates new indexing jobs for each connector / credential pair which is:
+    1. Enabled
+    2. `refresh_frequency` time has passed since the last indexing run for this pair
+    3. There is not already an ongoing indexing attempt for this pair
+    """
+    with Session(get_sqlalchemy_engine()) as db_session:
+        ongoing: set[tuple[int | None, int]] = set()
+        for attempt_id in existing_jobs:
+            attempt = get_index_attempt(
+                db_session=db_session, index_attempt_id=attempt_id
+            )
+            if attempt is None:
+                logger.error(
+                    f"Unable to find IndexAttempt for ID '{attempt_id}' when creating "
+                    "indexing jobs"
+                )
+                continue
+            ongoing.add(
+                (
+                    attempt.connector_credential_pair_id,
+                    attempt.search_settings_id,
+                )
+            )
+
+        # Get the primary search settings
+        primary_search_settings = get_current_search_settings(db_session)
+        search_settings = [primary_search_settings]
+
+        # Check for secondary search settings
+        secondary_search_settings = get_secondary_search_settings(db_session)
+        if secondary_search_settings is not None:
+            # If secondary settings exist, add them to the list
+            search_settings.append(secondary_search_settings)
+
+        all_connector_credential_pairs = fetch_connector_credential_pairs(db_session)
+        for cc_pair in all_connector_credential_pairs:
+            for search_settings_instance in search_settings:
+                # Check if there is an ongoing indexing attempt for this connector credential pair
+                if (cc_pair.id, search_settings_instance.id) in ongoing:
+                    continue
+
+                last_attempt = get_last_attempt_for_cc_pair(
+                    cc_pair.id, search_settings_instance.id, db_session
+                )
+                if not _should_create_new_indexing(
+                    cc_pair=cc_pair,
+                    last_index=last_attempt,
+                    search_settings_instance=search_settings_instance,
+                    secondary_index_building=len(search_settings) > 1,
+                    db_session=db_session,
+                ):
+                    continue
+
+                create_index_attempt(
+                    cc_pair.id, search_settings_instance.id, db_session
+                )
+
+
+def cleanup_indexing_jobs(
+    existing_jobs: dict[int, Future | SimpleJob],
+    timeout_hours: int = CLEANUP_INDEXING_JOBS_TIMEOUT,
+) -> dict[int, Future | SimpleJob]:
+    existing_jobs_copy = existing_jobs.copy()
+    # clean up completed jobs
+    with Session(get_sqlalchemy_engine()) as db_session:
+        for attempt_id, job in existing_jobs.items():
+            index_attempt = get_index_attempt(
+                db_session=db_session, index_attempt_id=attempt_id
+            )
+
+            # do nothing for ongoing jobs that haven't been stopped
+            if not job.done():
+                if not index_attempt:
+                    continue
+
+                if not index_attempt.is_finished():
+                    continue
+
+            if job.status == "error":
+                logger.error(job.exception())
+
+            job.release()
+            del existing_jobs_copy[attempt_id]
+
+            if not index_attempt:
+                logger.error(
+                    f"Unable to find IndexAttempt for ID '{attempt_id}' when cleaning "
+                    "up indexing jobs"
+                )
+                continue
+
+            if (
+                index_attempt.status == IndexingStatus.IN_PROGRESS
+                or job.status == "error"
+            ):
+                _mark_run_failed(
+                    db_session=db_session,
+                    index_attempt=index_attempt,
+                    failure_reason=_UNEXPECTED_STATE_FAILURE_REASON,
+                )
+
+        # clean up in-progress jobs that were never completed
+        connectors = fetch_connectors(db_session)
+        for connector in connectors:
+            in_progress_indexing_attempts = get_inprogress_index_attempts(
+                connector.id, db_session
+            )
+            for index_attempt in in_progress_indexing_attempts:
+                if index_attempt.id in existing_jobs:
+                    # If index attempt is canceled, stop the run
+                    if index_attempt.status == IndexingStatus.FAILED:
+                        existing_jobs[index_attempt.id].cancel()
+                    # check to see if the job has been updated in last `timeout_hours` hours, if not
+                    # assume it to frozen in some bad state and just mark it as failed. Note: this relies
+                    # on the fact that the `time_updated` field is constantly updated every
+                    # batch of documents indexed
+                    current_db_time = get_db_current_time(db_session=db_session)
+                    time_since_update = current_db_time - index_attempt.time_updated
+                    if time_since_update.total_seconds() > 60 * 60 * timeout_hours:
+                        existing_jobs[index_attempt.id].cancel()
+                        _mark_run_failed(
+                            db_session=db_session,
+                            index_attempt=index_attempt,
+                            failure_reason="Indexing run frozen - no updates in the last three hours. "
+                            "The run will be re-attempted at next scheduled indexing time.",
+                        )
+                else:
+                    # If job isn't known, simply mark it as failed
+                    _mark_run_failed(
+                        db_session=db_session,
+                        index_attempt=index_attempt,
+                        failure_reason=_UNEXPECTED_STATE_FAILURE_REASON,
+                    )
+
+    return existing_jobs_copy
+
+
+def kickoff_indexing_jobs(
+    existing_jobs: dict[int, Future | SimpleJob],
+    client: Client | SimpleJobClient,
+    secondary_client: Client | SimpleJobClient,
+) -> dict[int, Future | SimpleJob]:
+    existing_jobs_copy = existing_jobs.copy()
+    engine = get_sqlalchemy_engine()
+
+    # Don't include jobs waiting in the Dask queue that just haven't started running
+    # Also (rarely) don't include for jobs that started but haven't updated the indexing tables yet
+    with Session(engine) as db_session:
+        # get_not_started_index_attempts orders its returned results from oldest to newest
+        # we must process attempts in a FIFO manner to prevent connector starvation
+        new_indexing_attempts = [
+            (attempt, attempt.search_settings)
+            for attempt in get_not_started_index_attempts(db_session)
+            if attempt.id not in existing_jobs
+        ]
+
+    logger.debug(f"Found {len(new_indexing_attempts)} new indexing task(s).")
+
+    if not new_indexing_attempts:
+        return existing_jobs
+
+    indexing_attempt_count = 0
+
+    primary_client_full = False
+    secondary_client_full = False
+    for attempt, search_settings in new_indexing_attempts:
+        if primary_client_full and secondary_client_full:
+            break
+
+        use_secondary_index = (
+            search_settings.status == IndexModelStatus.FUTURE
+            if search_settings is not None
+            else False
+        )
+        if attempt.connector_credential_pair.connector is None:
+            logger.warning(
+                f"Skipping index attempt as Connector has been deleted: {attempt}"
+            )
+            with Session(engine) as db_session:
+                mark_attempt_failed(
+                    attempt, db_session, failure_reason="Connector is null"
+                )
+            continue
+        if attempt.connector_credential_pair.credential is None:
+            logger.warning(
+                f"Skipping index attempt as Credential has been deleted: {attempt}"
+            )
+            with Session(engine) as db_session:
+                mark_attempt_failed(
+                    attempt, db_session, failure_reason="Credential is null"
+                )
+            continue
+
+        if not use_secondary_index:
+            if not primary_client_full:
+                run = client.submit(
+                    run_indexing_entrypoint,
+                    attempt.id,
+                    attempt.connector_credential_pair_id,
+                    global_version.get_is_ee_version(),
+                    pure=False,
+                )
+                if not run:
+                    primary_client_full = True
+        else:
+            if not secondary_client_full:
+                run = secondary_client.submit(
+                    run_indexing_entrypoint,
+                    attempt.id,
+                    attempt.connector_credential_pair_id,
+                    global_version.get_is_ee_version(),
+                    pure=False,
+                )
+                if not run:
+                    secondary_client_full = True
+
+        if run:
+            if indexing_attempt_count == 0:
+                logger.info(
+                    f"Indexing dispatch starts: pending={len(new_indexing_attempts)}"
+                )
+
+            indexing_attempt_count += 1
+            secondary_str = " (secondary index)" if use_secondary_index else ""
+            logger.info(
+                f"Indexing dispatched{secondary_str}: "
+                f"attempt_id={attempt.id} "
+                f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
+                f"credentials='{attempt.connector_credential_pair.credential_id}'"
+            )
+            existing_jobs_copy[attempt.id] = run
+
+    if indexing_attempt_count > 0:
+        logger.info(
+            f"Indexing dispatch results: "
+            f"initial_pending={len(new_indexing_attempts)} "
+            f"started={indexing_attempt_count} "
+            f"remaining={len(new_indexing_attempts) - indexing_attempt_count}"
+        )
+
+    return existing_jobs_copy
+
+
+def update_loop(
+    delay: int = 10,
+    num_workers: int = NUM_INDEXING_WORKERS,
+    num_secondary_workers: int = NUM_SECONDARY_INDEXING_WORKERS,
+) -> None:
+    engine = get_sqlalchemy_engine()
+    with Session(engine) as db_session:
+        check_index_swap(db_session=db_session)
+        search_settings = get_current_search_settings(db_session)
+
+        # So that the first time users aren't surprised by really slow speed of first
+        # batch of documents indexed
+
+        if search_settings.provider_type is None:
+            logger.notice("Running a first inference to warm up embedding model")
+            embedding_model = EmbeddingModel.from_db_model(
+                search_settings=search_settings,
+                server_host=INDEXING_MODEL_SERVER_HOST,
+                server_port=MODEL_SERVER_PORT,
+            )
+
+            warm_up_bi_encoder(
+                embedding_model=embedding_model,
+            )
+            logger.notice("First inference complete.")
+
+    client_primary: Client | SimpleJobClient
+    client_secondary: Client | SimpleJobClient
+    if DASK_JOB_CLIENT_ENABLED:
+        cluster_primary = LocalCluster(
+            n_workers=num_workers,
+            threads_per_worker=1,
+            # there are warning about high memory usage + "Event loop unresponsive"
+            # which are not relevant to us since our workers are expected to use a
+            # lot of memory + involve CPU intensive tasks that will not relinquish
+            # the event loop
+            silence_logs=logging.ERROR,
+        )
+        cluster_secondary = LocalCluster(
+            n_workers=num_secondary_workers,
+            threads_per_worker=1,
+            silence_logs=logging.ERROR,
+        )
+        client_primary = Client(cluster_primary)
+        client_secondary = Client(cluster_secondary)
+        if LOG_LEVEL.lower() == "debug":
+            client_primary.register_worker_plugin(ResourceLogger())
+    else:
+        client_primary = SimpleJobClient(n_workers=num_workers)
+        client_secondary = SimpleJobClient(n_workers=num_secondary_workers)
+
+    existing_jobs: dict[int, Future | SimpleJob] = {}
+
+    logger.notice("Startup complete. Waiting for indexing jobs...")
+    while True:
+        start = time.time()
+        start_time_utc = datetime.utcfromtimestamp(start).strftime("%Y-%m-%d %H:%M:%S")
+        logger.debug(f"Running update, current UTC time: {start_time_utc}")
+
+        if existing_jobs:
+            # TODO: make this debug level once the "no jobs are being scheduled" issue is resolved
+            logger.debug(
+                "Found existing indexing jobs: "
+                f"{[(attempt_id, job.status) for attempt_id, job in existing_jobs.items()]}"
+            )
+
+        try:
+            with Session(get_sqlalchemy_engine()) as db_session:
+                check_index_swap(db_session)
+            existing_jobs = cleanup_indexing_jobs(existing_jobs=existing_jobs)
+            create_indexing_jobs(existing_jobs=existing_jobs)
+            existing_jobs = kickoff_indexing_jobs(
+                existing_jobs=existing_jobs,
+                client=client_primary,
+                secondary_client=client_secondary,
+            )
+        except Exception as e:
+            logger.exception(f"Failed to run update due to {e}")
+        sleep_time = delay - (time.time() - start)
+        if sleep_time > 0:
+            time.sleep(sleep_time)
+
+
+def update__main() -> None:
+    set_is_ee_based_on_env_variable()
+
+    # initialize the Postgres connection pool
+    SqlEngine.set_app_name(POSTGRES_INDEXER_APP_NAME)
+
+    logger.notice("Starting indexing service")
+    update_loop()
+
+
+if __name__ == "__main__":
+    update__main()

backend/danswer/chat/chat_utils.py

@@ -1,8 +1,6 @@
 import re
 from typing import cast
-from uuid import UUID
 
-from fastapi.datastructures import Headers
 from sqlalchemy.orm import Session
 
 from danswer.chat.models import CitationInfo
@@ -35,7 +33,7 @@ def llm_doc_from_inference_section(inference_section: InferenceSection) -> LlmDo
 
 
 def create_chat_chain(
-    chat_session_id: UUID,
+    chat_session_id: int,
     db_session: Session,
     prefetch_tool_calls: bool = True,
     # Optional id at which we finish processing
@@ -168,31 +166,3 @@ def reorganize_citations(
             new_citation_info[citation.citation_num] = citation
 
     return new_answer, list(new_citation_info.values())
-
-
-def extract_headers(
-    headers: dict[str, str] | Headers, pass_through_headers: list[str] | None
-) -> dict[str, str]:
-    """
-    Extract headers specified in pass_through_headers from input headers.
-    Handles both dict and FastAPI Headers objects, accounting for lowercase keys.
-
-    Args:
-        headers: Input headers as dict or Headers object.
-
-    Returns:
-        dict: Filtered headers based on pass_through_headers.
-    """
-    if not pass_through_headers:
-        return {}
-
-    extracted_headers: dict[str, str] = {}
-    for key in pass_through_headers:
-        if key in headers:
-            extracted_headers[key] = headers[key]
-        else:
-            # fastapi makes all header keys lowercase, handling that here
-            lowercase_key = key.lower()
-            if lowercase_key in headers:
-                extracted_headers[lowercase_key] = headers[lowercase_key]
-    return extracted_headers

backend/danswer/chat/process_message.py

@@ -18,10 +18,6 @@ from danswer.chat.models import MessageResponseIDInfo
 from danswer.chat.models import MessageSpecificCitations
 from danswer.chat.models import QADocsResponse
 from danswer.chat.models import StreamingError
-from danswer.configs.app_configs import AZURE_DALLE_API_BASE
-from danswer.configs.app_configs import AZURE_DALLE_API_KEY
-from danswer.configs.app_configs import AZURE_DALLE_API_VERSION
-from danswer.configs.app_configs import AZURE_DALLE_DEPLOYMENT_NAME
 from danswer.configs.chat_configs import BING_API_KEY
 from danswer.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from danswer.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
@@ -105,7 +101,6 @@ from danswer.tools.tool import ToolResponse
 from danswer.tools.tool_runner import ToolCallFinalResult
 from danswer.tools.utils import compute_all_tool_tokens
 from danswer.tools.utils import explicit_tool_calling_supported
-from danswer.utils.headers import header_dict_to_header_list
 from danswer.utils.logger import setup_logger
 from danswer.utils.timing import log_generator_function_time
 
@@ -277,7 +272,6 @@ def stream_chat_message_objects(
     # on the `new_msg_req.message`. Currently, requires a state where the last message is a
     use_existing_user_message: bool = False,
     litellm_additional_headers: dict[str, str] | None = None,
-    custom_tool_additional_headers: dict[str, str] | None = None,
     is_connected: Callable[[], bool] | None = None,
     enforce_chat_session_id_for_search_docs: bool = True,
 ) -> ChatPacketStream:
@@ -566,26 +560,7 @@ def stream_chat_message_objects(
                         and llm.config.api_key
                         and llm.config.model_provider == "openai"
                     ):
-                        img_generation_llm_config = LLMConfig(
-                            model_provider=llm.config.model_provider,
-                            model_name="dall-e-3",
-                            temperature=GEN_AI_TEMPERATURE,
-                            api_key=llm.config.api_key,
-                            api_base=llm.config.api_base,
-                            api_version=llm.config.api_version,
-                        )
-                    elif (
-                        llm.config.model_provider == "azure"
-                        and AZURE_DALLE_API_KEY is not None
-                    ):
-                        img_generation_llm_config = LLMConfig(
-                            model_provider="azure",
-                            model_name=f"azure/{AZURE_DALLE_DEPLOYMENT_NAME}",
-                            temperature=GEN_AI_TEMPERATURE,
-                            api_key=AZURE_DALLE_API_KEY,
-                            api_base=AZURE_DALLE_API_BASE,
-                            api_version=AZURE_DALLE_API_VERSION,
-                        )
+                        img_generation_llm_config = llm.config
                     else:
                         llm_providers = fetch_existing_llm_providers(db_session)
                         openai_provider = next(
@@ -604,7 +579,7 @@ def stream_chat_message_objects(
                             )
                         img_generation_llm_config = LLMConfig(
                             model_provider=openai_provider.provider,
-                            model_name="dall-e-3",
+                            model_name=openai_provider.default_model_name,
                             temperature=GEN_AI_TEMPERATURE,
                             api_key=openai_provider.api_key,
                             api_base=openai_provider.api_base,
@@ -616,7 +591,6 @@ def stream_chat_message_objects(
                             api_base=img_generation_llm_config.api_base,
                             api_version=img_generation_llm_config.api_version,
                             additional_headers=litellm_additional_headers,
-                            model=img_generation_llm_config.model_name,
                         )
                     ]
                 elif tool_cls.__name__ == InternetSearchTool.__name__:
@@ -641,12 +615,7 @@ def stream_chat_message_objects(
                             chat_session_id=chat_session_id,
                             message_id=user_message.id if user_message else None,
                         ),
-                        custom_headers=(db_tool_model.custom_headers or [])
-                        + (
-                            header_dict_to_header_list(
-                                custom_tool_additional_headers or {}
-                            )
-                        ),
+                        custom_headers=db_tool_model.custom_headers,
                     ),
                 )
 
@@ -672,7 +641,6 @@ def stream_chat_message_objects(
                     all_docs_useful=selected_db_search_docs is not None
                 ),
                 document_pruning_config=document_pruning_config,
-                structured_response_format=new_msg_req.structured_response_format,
             ),
             prompt_config=prompt_config,
             llm=(
@@ -870,7 +838,6 @@ def stream_chat_message(
     user: User | None,
     use_existing_user_message: bool = False,
     litellm_additional_headers: dict[str, str] | None = None,
-    custom_tool_additional_headers: dict[str, str] | None = None,
     is_connected: Callable[[], bool] | None = None,
 ) -> Iterator[str]:
     with get_session_context_manager() as db_session:
@@ -880,7 +847,6 @@ def stream_chat_message(
             db_session=db_session,
             use_existing_user_message=use_existing_user_message,
             litellm_additional_headers=litellm_additional_headers,
-            custom_tool_additional_headers=custom_tool_additional_headers,
             is_connected=is_connected,
         )
         for obj in objects:

backend/danswer/configs/app_configs.py

@@ -43,9 +43,6 @@ WEB_DOMAIN = os.environ.get("WEB_DOMAIN") or "http://localhost:3000"
 AUTH_TYPE = AuthType((os.environ.get("AUTH_TYPE") or AuthType.DISABLED.value).lower())
 DISABLE_AUTH = AUTH_TYPE == AuthType.DISABLED
 
-# Necessary for cloud integration tests
-DISABLE_VERIFICATION = os.environ.get("DISABLE_VERIFICATION", "").lower() == "true"
-
 # Encryption key secret is used to encrypt connector credentials, api keys, and other sensitive
 # information. This provides an extra layer of security on top of Postgres access controls
 # and is available in Danswer EE
@@ -56,6 +53,7 @@ MASK_CREDENTIAL_PREFIX = (
     os.environ.get("MASK_CREDENTIAL_PREFIX", "True").lower() != "false"
 )
 
+
 SESSION_EXPIRE_TIME_SECONDS = int(
     os.environ.get("SESSION_EXPIRE_TIME_SECONDS") or 86400 * 7
 )  # 7 days
@@ -118,22 +116,17 @@ VESPA_HOST = os.environ.get("VESPA_HOST") or "localhost"
 VESPA_CONFIG_SERVER_HOST = os.environ.get("VESPA_CONFIG_SERVER_HOST") or VESPA_HOST
 VESPA_PORT = os.environ.get("VESPA_PORT") or "8081"
 VESPA_TENANT_PORT = os.environ.get("VESPA_TENANT_PORT") or "19071"
-
-VESPA_CLOUD_URL = os.environ.get("VESPA_CLOUD_URL", "")
-
 # The default below is for dockerized deployment
 VESPA_DEPLOYMENT_ZIP = (
     os.environ.get("VESPA_DEPLOYMENT_ZIP") or "/app/danswer/vespa-app.zip"
 )
-VESPA_CLOUD_CERT_PATH = os.environ.get("VESPA_CLOUD_CERT_PATH")
-VESPA_CLOUD_KEY_PATH = os.environ.get("VESPA_CLOUD_KEY_PATH")
-
 # Number of documents in a batch during indexing (further batching done by chunks before passing to bi-encoder)
 try:
     INDEX_BATCH_SIZE = int(os.environ.get("INDEX_BATCH_SIZE", 16))
 except ValueError:
     INDEX_BATCH_SIZE = 16
 
+
 # Below are intended to match the env variables names used by the official postgres docker image
 # https://hub.docker.com/_/postgres
 POSTGRES_USER = os.environ.get("POSTGRES_USER") or "postgres"
@@ -200,41 +193,6 @@ try:
 except ValueError:
     CELERY_BROKER_POOL_LIMIT = CELERY_BROKER_POOL_LIMIT_DEFAULT
 
-CELERY_WORKER_LIGHT_CONCURRENCY_DEFAULT = 24
-try:
-    CELERY_WORKER_LIGHT_CONCURRENCY = int(
-        os.environ.get(
-            "CELERY_WORKER_LIGHT_CONCURRENCY", CELERY_WORKER_LIGHT_CONCURRENCY_DEFAULT
-        )
-    )
-except ValueError:
-    CELERY_WORKER_LIGHT_CONCURRENCY = CELERY_WORKER_LIGHT_CONCURRENCY_DEFAULT
-
-CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER_DEFAULT = 8
-try:
-    CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER = int(
-        os.environ.get(
-            "CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER",
-            CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER_DEFAULT,
-        )
-    )
-except ValueError:
-    CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER = (
-        CELERY_WORKER_LIGHT_PREFETCH_MULTIPLIER_DEFAULT
-    )
-
-CELERY_WORKER_INDEXING_CONCURRENCY_DEFAULT = 1
-try:
-    env_value = os.environ.get("CELERY_WORKER_INDEXING_CONCURRENCY")
-    if not env_value:
-        env_value = os.environ.get("NUM_INDEXING_WORKERS")
-
-    if not env_value:
-        env_value = str(CELERY_WORKER_INDEXING_CONCURRENCY_DEFAULT)
-    CELERY_WORKER_INDEXING_CONCURRENCY = int(env_value)
-except ValueError:
-    CELERY_WORKER_INDEXING_CONCURRENCY = CELERY_WORKER_INDEXING_CONCURRENCY_DEFAULT
-
 #####
 # Connector Configs
 #####
@@ -290,6 +248,12 @@ CONFLUENCE_CONNECTOR_INDEX_ARCHIVED_PAGES = (
     os.environ.get("CONFLUENCE_CONNECTOR_INDEX_ARCHIVED_PAGES", "").lower() == "true"
 )
 
+# Save pages labels as Danswer metadata tags
+# The reason to skip this would be to reduce the number of calls to Confluence due to rate limit concerns
+CONFLUENCE_CONNECTOR_SKIP_LABEL_INDEXING = (
+    os.environ.get("CONFLUENCE_CONNECTOR_SKIP_LABEL_INDEXING", "").lower() == "true"
+)
+
 # Attachments exceeding this size will not be retrieved (in bytes)
 CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD = int(
     os.environ.get("CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD", 10 * 1024 * 1024)
@@ -437,11 +401,6 @@ CUSTOM_ANSWER_VALIDITY_CONDITIONS = json.loads(
     os.environ.get("CUSTOM_ANSWER_VALIDITY_CONDITIONS", "[]")
 )
 
-VESPA_REQUEST_TIMEOUT = int(os.environ.get("VESPA_REQUEST_TIMEOUT") or "5")
-
-SYSTEM_RECURSION_LIMIT = int(os.environ.get("SYSTEM_RECURSION_LIMIT") or "1000")
-
-PARSE_WITH_TRAFILATURA = os.environ.get("PARSE_WITH_TRAFILATURA", "").lower() == "true"
 
 #####
 # Enterprise Edition Configs
@@ -454,38 +413,10 @@ ENTERPRISE_EDITION_ENABLED = (
     os.environ.get("ENABLE_PAID_ENTERPRISE_EDITION_FEATURES", "").lower() == "true"
 )
 
-# Azure DALL-E Configurations
-AZURE_DALLE_API_VERSION = os.environ.get("AZURE_DALLE_API_VERSION")
-AZURE_DALLE_API_KEY = os.environ.get("AZURE_DALLE_API_KEY")
-AZURE_DALLE_API_BASE = os.environ.get("AZURE_DALLE_API_BASE")
-AZURE_DALLE_DEPLOYMENT_NAME = os.environ.get("AZURE_DALLE_DEPLOYMENT_NAME")
 
-
-# Cloud configuration
-
-# Multi-tenancy configuration
 MULTI_TENANT = os.environ.get("MULTI_TENANT", "").lower() == "true"
+SECRET_JWT_KEY = os.environ.get("SECRET_JWT_KEY", "")
 
-# Use managed Vespa (Vespa Cloud). If set, must also set VESPA_CLOUD_URL, VESPA_CLOUD_CERT_PATH and VESPA_CLOUD_KEY_PATH
-MANAGED_VESPA = os.environ.get("MANAGED_VESPA", "").lower() == "true"
 
-ENABLE_EMAIL_INVITES = os.environ.get("ENABLE_EMAIL_INVITES", "").lower() == "true"
-
-# Security and authentication
-SECRET_JWT_KEY = os.environ.get(
-    "SECRET_JWT_KEY", ""
-)  # Used for encryption of the JWT token for user's tenant context
-DATA_PLANE_SECRET = os.environ.get(
-    "DATA_PLANE_SECRET", ""
-)  # Used for secure communication between the control and data plane
-EXPECTED_API_KEY = os.environ.get(
-    "EXPECTED_API_KEY", ""
-)  # Additional security check for the control plane API
-
-# API configuration
-CONTROL_PLANE_API_BASE_URL = os.environ.get(
-    "CONTROL_PLANE_API_BASE_URL", "http://localhost:8082"
-)
-
-# JWT configuration
-JWT_ALGORITHM = "HS256"
+DATA_PLANE_SECRET = os.environ.get("DATA_PLANE_SECRET", "")
+EXPECTED_API_KEY = os.environ.get("EXPECTED_API_KEY", "")

backend/danswer/configs/constants.py

@@ -31,9 +31,6 @@ DISABLED_GEN_AI_MSG = (
     "You can still use Danswer as a search engine."
 )
 
-# Prefix used for all tenant ids
-TENANT_ID_PREFIX = "tenant_"
-
 # Postgres connection constants for application_name
 POSTGRES_WEB_APP_NAME = "web"
 POSTGRES_INDEXER_APP_NAME = "indexer"
@@ -42,10 +39,9 @@ POSTGRES_CELERY_BEAT_APP_NAME = "celery_beat"
 POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME = "celery_worker_primary"
 POSTGRES_CELERY_WORKER_LIGHT_APP_NAME = "celery_worker_light"
 POSTGRES_CELERY_WORKER_HEAVY_APP_NAME = "celery_worker_heavy"
-POSTGRES_CELERY_WORKER_INDEXING_APP_NAME = "celery_worker_indexing"
-POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME = "celery_worker_indexing_child"
 POSTGRES_PERMISSIONS_APP_NAME = "permissions"
 POSTGRES_UNKNOWN_APP_NAME = "unknown"
+POSTGRES_DEFAULT_SCHEMA = "public"
 
 # API Keys
 DANSWER_API_KEY_PREFIX = "API_KEY__"
@@ -70,21 +66,10 @@ KV_CUSTOMER_UUID_KEY = "customer_uuid"
 KV_INSTANCE_DOMAIN_KEY = "instance_domain"
 KV_ENTERPRISE_SETTINGS_KEY = "danswer_enterprise_settings"
 KV_CUSTOM_ANALYTICS_SCRIPT_KEY = "__custom_analytics_script__"
-KV_DOCUMENTS_SEEDED_KEY = "documents_seeded"
 
 CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT = 60
 CELERY_PRIMARY_WORKER_LOCK_TIMEOUT = 120
 
-# needs to be long enough to cover the maximum time it takes to download an object
-# if we can get callbacks as object bytes download, we could lower this a lot.
-CELERY_INDEXING_LOCK_TIMEOUT = 60 * 60  # 60 min
-
-# needs to be long enough to cover the maximum time it takes to download an object
-# if we can get callbacks as object bytes download, we could lower this a lot.
-CELERY_PRUNING_LOCK_TIMEOUT = 300  # 5 min
-
-DANSWER_REDIS_FUNCTION_LOCK_PREFIX = "da_function_lock:"
-
 
 class DocumentSource(str, Enum):
     # Special case, document passed in via Danswer APIs without specifying a source type
@@ -130,13 +115,8 @@ class DocumentSource(str, Enum):
     NOT_APPLICABLE = "not_applicable"
 
 
-DocumentSourceRequiringTenantContext: list[DocumentSource] = [DocumentSource.FILE]
-
-
 class NotificationType(str, Enum):
     REINDEX = "reindex"
-    PERSONA_SHARED = "persona_shared"
-    TRIAL_ENDS_TWO_DAYS = "two_day_trial_ending"  # 2 days left in trial
 
 
 class BlobType(str, Enum):
@@ -161,9 +141,6 @@ class AuthType(str, Enum):
     OIDC = "oidc"
     SAML = "saml"
 
-    # google auth and basic
-    CLOUD = "cloud"
-
 
 class SessionType(str, Enum):
     CHAT = "Chat"
@@ -213,19 +190,14 @@ class DanswerCeleryQueues:
     VESPA_METADATA_SYNC = "vespa_metadata_sync"
     CONNECTOR_DELETION = "connector_deletion"
     CONNECTOR_PRUNING = "connector_pruning"
-    CONNECTOR_INDEXING = "connector_indexing"
 
 
 class DanswerRedisLocks:
     PRIMARY_WORKER = "da_lock:primary_worker"
     CHECK_VESPA_SYNC_BEAT_LOCK = "da_lock:check_vespa_sync_beat"
+    MONITOR_VESPA_SYNC_BEAT_LOCK = "da_lock:monitor_vespa_sync_beat"
     CHECK_CONNECTOR_DELETION_BEAT_LOCK = "da_lock:check_connector_deletion_beat"
     CHECK_PRUNE_BEAT_LOCK = "da_lock:check_prune_beat"
-    CHECK_INDEXING_BEAT_LOCK = "da_lock:check_indexing_beat"
-    MONITOR_VESPA_SYNC_BEAT_LOCK = "da_lock:monitor_vespa_sync_beat"
-
-    PRUNING_LOCK_PREFIX = "da_lock:pruning"
-    INDEXING_METADATA_PREFIX = "da_metadata:indexing"
 
 
 class DanswerCeleryPriority(int, Enum):

backend/danswer/configs/tool_configs.py

@@ -1,22 +0,0 @@
-import json
-import os
-
-
-# if specified, will pass through request headers to the call to API calls made by custom tools
-CUSTOM_TOOL_PASS_THROUGH_HEADERS: list[str] | None = None
-_CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW = os.environ.get(
-    "CUSTOM_TOOL_PASS_THROUGH_HEADERS"
-)
-if _CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW:
-    try:
-        CUSTOM_TOOL_PASS_THROUGH_HEADERS = json.loads(
-            _CUSTOM_TOOL_PASS_THROUGH_HEADERS_RAW
-        )
-    except Exception:
-        # need to import here to avoid circular imports
-        from danswer.utils.logger import setup_logger
-
-        logger = setup_logger()
-        logger.error(
-            "Failed to parse CUSTOM_TOOL_PASS_THROUGH_HEADERS, must be a valid JSON object"
-        )

backend/danswer/connectors/README.md

@@ -13,8 +13,8 @@ Connectors come in 3 different flows:
   documents via a connector's API or loads the documents from some sort of a dump file.
 - Poll connector:
   - Incrementally updates documents based on a provided time range. It is used by the background job to pull the latest
-  changes and additions since the last round of polling. This connector helps keep the document index up to date
-  without needing to fetch/embed/index every document which would be too slow to do frequently on large sets of
+  changes additions and changes since the last round of polling. This connector helps keep the document index up to date
+  without needing to fetch/embed/index every document which generally be too slow to do frequently on large sets of
   documents.
 - Event Based connectors:
   - Connectors that listen to events and update documents accordingly.

backend/danswer/connectors/axero/connector.py

@@ -15,6 +15,7 @@ from danswer.connectors.cross_connector_utils.miscellaneous_utils import time_st
 from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
     rate_limit_builder,
 )
+from danswer.connectors.cross_connector_utils.retry_wrapper import retry_builder
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import PollConnector
 from danswer.connectors.interfaces import SecondsSinceUnixEpoch
@@ -23,7 +24,6 @@ from danswer.connectors.models import Document
 from danswer.connectors.models import Section
 from danswer.file_processing.html_utils import parse_html_page_basic
 from danswer.utils.logger import setup_logger
-from danswer.utils.retry_wrapper import retry_builder
 
 
 logger = setup_logger()

backend/danswer/connectors/clickup/connector.py

@@ -10,6 +10,7 @@ from danswer.configs.constants import DocumentSource
 from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
     rate_limit_builder,
 )
+from danswer.connectors.cross_connector_utils.retry_wrapper import retry_builder
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import LoadConnector
 from danswer.connectors.interfaces import PollConnector
@@ -18,7 +19,6 @@ from danswer.connectors.models import BasicExpertInfo
 from danswer.connectors.models import ConnectorMissingCredentialError
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
-from danswer.utils.retry_wrapper import retry_builder
 
 
 CLICKUP_API_BASE_URL = "https://api.clickup.com/api/v2"

backend/danswer/connectors/confluence/confluence_utils.py

@@ -0,0 +1,32 @@
+import bs4
+
+
+def build_confluence_document_id(base_url: str, content_url: str) -> str:
+    """For confluence, the document id is the page url for a page based document
+        or the attachment download url for an attachment based document
+
+    Args:
+        base_url (str): The base url of the Confluence instance
+        content_url (str): The url of the page or attachment download url
+
+    Returns:
+        str: The document id
+    """
+    return f"{base_url}{content_url}"
+
+
+def get_used_attachments(text: str) -> list[str]:
+    """Parse a Confluence html page to generate a list of current
+        attachment in used
+
+    Args:
+        text (str): The page content
+
+    Returns:
+        list[str]: List of filenames currently in use by the page text
+    """
+    files_in_used = []
+    soup = bs4.BeautifulSoup(text, "html.parser")
+    for attachment in soup.findAll("ri:attachment"):
+        files_in_used.append(attachment.attrs["ri:filename"])
+    return files_in_used

backend/danswer/connectors/confluence/onyx_confluence.py

@@ -1,226 +0,0 @@
-import math
-import time
-from collections.abc import Callable
-from collections.abc import Iterator
-from typing import Any
-from typing import cast
-from typing import TypeVar
-from urllib.parse import quote
-
-from atlassian import Confluence  # type:ignore
-from requests import HTTPError
-
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-F = TypeVar("F", bound=Callable[..., Any])
-
-
-RATE_LIMIT_MESSAGE_LOWERCASE = "Rate limit exceeded".lower()
-
-
-class ConfluenceRateLimitError(Exception):
-    pass
-
-
-def _handle_http_error(e: HTTPError, attempt: int) -> int:
-    MIN_DELAY = 2
-    MAX_DELAY = 60
-    STARTING_DELAY = 5
-    BACKOFF = 2
-
-    # Check if the response or headers are None to avoid potential AttributeError
-    if e.response is None or e.response.headers is None:
-        logger.warning("HTTPError with `None` as response or as headers")
-        raise e
-
-    if (
-        e.response.status_code != 429
-        and RATE_LIMIT_MESSAGE_LOWERCASE not in e.response.text.lower()
-    ):
-        raise e
-
-    retry_after = None
-
-    retry_after_header = e.response.headers.get("Retry-After")
-    if retry_after_header is not None:
-        try:
-            retry_after = int(retry_after_header)
-            if retry_after > MAX_DELAY:
-                logger.warning(
-                    f"Clamping retry_after from {retry_after} to {MAX_DELAY} seconds..."
-                )
-                retry_after = MAX_DELAY
-            if retry_after < MIN_DELAY:
-                retry_after = MIN_DELAY
-        except ValueError:
-            pass
-
-    if retry_after is not None:
-        logger.warning(
-            f"Rate limiting with retry header. Retrying after {retry_after} seconds..."
-        )
-        delay = retry_after
-    else:
-        logger.warning(
-            "Rate limiting without retry header. Retrying with exponential backoff..."
-        )
-        delay = min(STARTING_DELAY * (BACKOFF**attempt), MAX_DELAY)
-
-    delay_until = math.ceil(time.monotonic() + delay)
-    return delay_until
-
-
-# https://developer.atlassian.com/cloud/confluence/rate-limiting/
-# this uses the native rate limiting option provided by the
-# confluence client and otherwise applies a simpler set of error handling
-def handle_confluence_rate_limit(confluence_call: F) -> F:
-    def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:
-        MAX_RETRIES = 5
-
-        TIMEOUT = 3600
-        timeout_at = time.monotonic() + TIMEOUT
-
-        for attempt in range(MAX_RETRIES):
-            if time.monotonic() > timeout_at:
-                raise TimeoutError(
-                    f"Confluence call attempts took longer than {TIMEOUT} seconds."
-                )
-
-            try:
-                # we're relying more on the client to rate limit itself
-                # and applying our own retries in a more specific set of circumstances
-                return confluence_call(*args, **kwargs)
-            except HTTPError as e:
-                delay_until = _handle_http_error(e, attempt)
-                while time.monotonic() < delay_until:
-                    # in the future, check a signal here to exit
-                    time.sleep(1)
-            except AttributeError as e:
-                # Some error within the Confluence library, unclear why it fails.
-                # Users reported it to be intermittent, so just retry
-                if attempt == MAX_RETRIES - 1:
-                    raise e
-
-                logger.exception(
-                    "Confluence Client raised an AttributeError. Retrying..."
-                )
-                time.sleep(5)
-
-    return cast(F, wrapped_call)
-
-
-_DEFAULT_PAGINATION_LIMIT = 100
-
-
-class OnyxConfluence(Confluence):
-    """
-    This is a custom Confluence class that overrides the default Confluence class to add a custom CQL method.
-    This is necessary because the default Confluence class does not properly support cql expansions.
-    All methods are automatically wrapped with handle_confluence_rate_limit.
-    """
-
-    def __init__(self, url: str, *args: Any, **kwargs: Any) -> None:
-        super(OnyxConfluence, self).__init__(url, *args, **kwargs)
-        self._wrap_methods()
-
-    def _wrap_methods(self) -> None:
-        """
-        For each attribute that is callable (i.e., a method) and doesn't start with an underscore,
-        wrap it with handle_confluence_rate_limit.
-        """
-        for attr_name in dir(self):
-            if callable(getattr(self, attr_name)) and not attr_name.startswith("_"):
-                setattr(
-                    self,
-                    attr_name,
-                    handle_confluence_rate_limit(getattr(self, attr_name)),
-                )
-
-    def _paginate_url(
-        self, url_suffix: str, limit: int | None = None
-    ) -> Iterator[list[dict[str, Any]]]:
-        """
-        This will paginate through the top level query.
-        """
-        if not limit:
-            limit = _DEFAULT_PAGINATION_LIMIT
-
-        connection_char = "&" if "?" in url_suffix else "?"
-        url_suffix += f"{connection_char}limit={limit}"
-
-        while url_suffix:
-            try:
-                next_response = self.get(url_suffix)
-            except Exception as e:
-                logger.exception("Error in danswer_cql: \n")
-                raise e
-            yield next_response.get("results", [])
-            url_suffix = next_response.get("_links", {}).get("next")
-
-    def paginated_groups_retrieval(
-        self,
-        limit: int | None = None,
-    ) -> Iterator[list[dict[str, Any]]]:
-        return self._paginate_url("rest/api/group", limit)
-
-    def paginated_group_members_retrieval(
-        self,
-        group_name: str,
-        limit: int | None = None,
-    ) -> Iterator[list[dict[str, Any]]]:
-        group_name = quote(group_name)
-        return self._paginate_url(f"rest/api/group/{group_name}/member", limit)
-
-    def paginated_cql_user_retrieval(
-        self,
-        cql: str,
-        expand: str | None = None,
-        limit: int | None = None,
-    ) -> Iterator[list[dict[str, Any]]]:
-        expand_string = f"&expand={expand}" if expand else ""
-        return self._paginate_url(
-            f"rest/api/search/user?cql={cql}{expand_string}", limit
-        )
-
-    def paginated_cql_page_retrieval(
-        self,
-        cql: str,
-        expand: str | None = None,
-        limit: int | None = None,
-    ) -> Iterator[list[dict[str, Any]]]:
-        expand_string = f"&expand={expand}" if expand else ""
-        return self._paginate_url(
-            f"rest/api/content/search?cql={cql}{expand_string}", limit
-        )
-
-    def cql_paginate_all_expansions(
-        self,
-        cql: str,
-        expand: str | None = None,
-        limit: int | None = None,
-    ) -> Iterator[list[dict[str, Any]]]:
-        """
-        This function will paginate through the top level query first, then
-        paginate through all of the expansions.
-        The limit only applies to the top level query.
-        All expansion paginations use default pagination limit (defined by Atlassian).
-        """
-
-        def _traverse_and_update(data: dict | list) -> None:
-            if isinstance(data, dict):
-                next_url = data.get("_links", {}).get("next")
-                if next_url and "results" in data:
-                    data["results"].extend(self._paginate_url(next_url))
-
-                for value in data.values():
-                    _traverse_and_update(value)
-            elif isinstance(data, list):
-                for item in data:
-                    _traverse_and_update(item)
-
-        for results in self.paginated_cql_page_retrieval(cql, expand, limit):
-            _traverse_and_update(results)
-            yield results

backend/danswer/connectors/confluence/rate_limit_handler.py

@@ -0,0 +1,82 @@
+import time
+from collections.abc import Callable
+from typing import Any
+from typing import cast
+from typing import TypeVar
+
+from requests import HTTPError
+
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+RATE_LIMIT_MESSAGE_LOWERCASE = "Rate limit exceeded".lower()
+
+
+class ConfluenceRateLimitError(Exception):
+    pass
+
+
+def make_confluence_call_handle_rate_limit(confluence_call: F) -> F:
+    def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:
+        max_retries = 5
+        starting_delay = 5
+        backoff = 2
+        max_delay = 600
+
+        for attempt in range(max_retries):
+            try:
+                return confluence_call(*args, **kwargs)
+            except HTTPError as e:
+                # Check if the response or headers are None to avoid potential AttributeError
+                if e.response is None or e.response.headers is None:
+                    logger.warning("HTTPError with `None` as response or as headers")
+                    raise e
+
+                retry_after_header = e.response.headers.get("Retry-After")
+                if (
+                    e.response.status_code == 429
+                    or RATE_LIMIT_MESSAGE_LOWERCASE in e.response.text.lower()
+                ):
+                    retry_after = None
+                    if retry_after_header is not None:
+                        try:
+                            retry_after = int(retry_after_header)
+                        except ValueError:
+                            pass
+
+                    if retry_after is not None:
+                        if retry_after > 600:
+                            logger.warning(
+                                f"Clamping retry_after from {retry_after} to {max_delay} seconds..."
+                            )
+                            retry_after = max_delay
+
+                        logger.warning(
+                            f"Rate limit hit. Retrying after {retry_after} seconds..."
+                        )
+                        time.sleep(retry_after)
+                    else:
+                        logger.warning(
+                            "Rate limit hit. Retrying with exponential backoff..."
+                        )
+                        delay = min(starting_delay * (backoff**attempt), max_delay)
+                        time.sleep(delay)
+                else:
+                    # re-raise, let caller handle
+                    raise
+            except AttributeError as e:
+                # Some error within the Confluence library, unclear why it fails.
+                # Users reported it to be intermittent, so just retry
+                logger.warning(f"Confluence Internal Error, retrying... {e}")
+                delay = min(starting_delay * (backoff**attempt), max_delay)
+                time.sleep(delay)
+
+                if attempt == max_retries - 1:
+                    raise e
+
+    return cast(F, wrapped_call)

backend/danswer/connectors/confluence/utils.py

@@ -1,214 +0,0 @@
-import io
-from datetime import datetime
-from datetime import timezone
-from typing import Any
-
-import bs4
-
-from danswer.configs.app_configs import (
-    CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD,
-)
-from danswer.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD
-from danswer.connectors.confluence.onyx_confluence import (
-    OnyxConfluence,
-)
-from danswer.file_processing.extract_file_text import extract_file_text
-from danswer.file_processing.html_utils import format_document_soup
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-_USER_EMAIL_CACHE: dict[str, str | None] = {}
-
-
-def get_user_email_from_username__server(
-    confluence_client: OnyxConfluence, user_name: str
-) -> str | None:
-    global _USER_EMAIL_CACHE
-    if _USER_EMAIL_CACHE.get(user_name) is None:
-        try:
-            response = confluence_client.get_mobile_parameters(user_name)
-            email = response.get("email")
-        except Exception:
-            email = None
-        _USER_EMAIL_CACHE[user_name] = email
-    return _USER_EMAIL_CACHE[user_name]
-
-
-_USER_NOT_FOUND = "Unknown Confluence User"
-_USER_ID_TO_DISPLAY_NAME_CACHE: dict[str, str | None] = {}
-
-
-def _get_user(confluence_client: OnyxConfluence, user_id: str) -> str:
-    """Get Confluence Display Name based on the account-id or userkey value
-
-    Args:
-        user_id (str): The user id (i.e: the account-id or userkey)
-        confluence_client (Confluence): The Confluence Client
-
-    Returns:
-        str: The User Display Name. 'Unknown User' if the user is deactivated or not found
-    """
-    global _USER_ID_TO_DISPLAY_NAME_CACHE
-    if _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) is None:
-        try:
-            result = confluence_client.get_user_details_by_userkey(user_id)
-            found_display_name = result.get("displayName")
-        except Exception:
-            found_display_name = None
-
-        if not found_display_name:
-            try:
-                result = confluence_client.get_user_details_by_accountid(user_id)
-                found_display_name = result.get("displayName")
-            except Exception:
-                found_display_name = None
-
-        _USER_ID_TO_DISPLAY_NAME_CACHE[user_id] = found_display_name
-
-    return _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) or _USER_NOT_FOUND
-
-
-def extract_text_from_confluence_html(
-    confluence_client: OnyxConfluence, confluence_object: dict[str, Any]
-) -> str:
-    """Parse a Confluence html page and replace the 'user Id' by the real
-        User Display Name
-
-    Args:
-        confluence_object (dict): The confluence object as a dict
-        confluence_client (Confluence): Confluence client
-
-    Returns:
-        str: loaded and formated Confluence page
-    """
-    body = confluence_object["body"]
-    object_html = body.get("storage", body.get("view", {})).get("value")
-
-    soup = bs4.BeautifulSoup(object_html, "html.parser")
-    for user in soup.findAll("ri:user"):
-        user_id = (
-            user.attrs["ri:account-id"]
-            if "ri:account-id" in user.attrs
-            else user.get("ri:userkey")
-        )
-        if not user_id:
-            logger.warning(
-                "ri:userkey not found in ri:user element. " f"Found attrs: {user.attrs}"
-            )
-            continue
-        # Include @ sign for tagging, more clear for LLM
-        user.replaceWith("@" + _get_user(confluence_client, user_id))
-    return format_document_soup(soup)
-
-
-def attachment_to_content(
-    confluence_client: OnyxConfluence,
-    attachment: dict[str, Any],
-) -> str | None:
-    """If it returns None, assume that we should skip this attachment."""
-    if attachment["metadata"]["mediaType"] in [
-        "image/jpeg",
-        "image/png",
-        "image/gif",
-        "image/svg+xml",
-        "video/mp4",
-        "video/quicktime",
-    ]:
-        return None
-
-    download_link = confluence_client.url + attachment["_links"]["download"]
-
-    attachment_size = attachment["extensions"]["fileSize"]
-    if attachment_size > CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD:
-        logger.warning(
-            f"Skipping {download_link} due to size. "
-            f"size={attachment_size} "
-            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD}"
-        )
-        return None
-
-    logger.info(f"_attachment_to_content - _session.get: link={download_link}")
-    response = confluence_client._session.get(download_link)
-    if response.status_code != 200:
-        logger.warning(
-            f"Failed to fetch {download_link} with invalid status code {response.status_code}"
-        )
-        return None
-
-    extracted_text = extract_file_text(
-        io.BytesIO(response.content),
-        file_name=attachment["title"],
-        break_on_unprocessable=False,
-    )
-    if len(extracted_text) > CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD:
-        logger.warning(
-            f"Skipping {download_link} due to char count. "
-            f"char count={len(extracted_text)} "
-            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD}"
-        )
-        return None
-
-    return extracted_text
-
-
-def build_confluence_document_id(base_url: str, content_url: str) -> str:
-    """For confluence, the document id is the page url for a page based document
-        or the attachment download url for an attachment based document
-
-    Args:
-        base_url (str): The base url of the Confluence instance
-        content_url (str): The url of the page or attachment download url
-
-    Returns:
-        str: The document id
-    """
-    return f"{base_url}{content_url}"
-
-
-def extract_referenced_attachment_names(page_text: str) -> list[str]:
-    """Parse a Confluence html page to generate a list of current
-        attachments in use
-
-    Args:
-        text (str): The page content
-
-    Returns:
-        list[str]: List of filenames currently in use by the page text
-    """
-    referenced_attachment_filenames = []
-    soup = bs4.BeautifulSoup(page_text, "html.parser")
-    for attachment in soup.findAll("ri:attachment"):
-        referenced_attachment_filenames.append(attachment.attrs["ri:filename"])
-    return referenced_attachment_filenames
-
-
-def datetime_from_string(datetime_string: str) -> datetime:
-    datetime_object = datetime.fromisoformat(datetime_string)
-
-    if datetime_object.tzinfo is None:
-        # If no timezone info, assume it is UTC
-        datetime_object = datetime_object.replace(tzinfo=timezone.utc)
-    else:
-        # If not in UTC, translate it
-        datetime_object = datetime_object.astimezone(timezone.utc)
-
-    return datetime_object
-
-
-def build_confluence_client(
-    credentials_json: dict[str, Any], is_cloud: bool, wiki_base: str
-) -> OnyxConfluence:
-    return OnyxConfluence(
-        api_version="cloud" if is_cloud else "latest",
-        # Remove trailing slash from wiki_base if present
-        url=wiki_base.rstrip("/"),
-        # passing in username causes issues for Confluence data center
-        username=credentials_json["confluence_username"] if is_cloud else None,
-        password=credentials_json["confluence_access_token"] if is_cloud else None,
-        token=credentials_json["confluence_access_token"] if not is_cloud else None,
-        backoff_and_retry=True,
-        max_backoff_retries=60,
-        max_backoff_seconds=60,
-    )

backend/danswer/connectors/cross_connector_utils/retry_wrapper.py

@@ -22,18 +22,18 @@ def retry_builder(
     jitter: tuple[float, float] | float = 1,
 ) -> Callable[[F], F]:
     """Builds a generic wrapper/decorator for calls to external APIs that
-    may fail due to rate limiting, flakes, or other reasons. Applies exponential
+    may fail due to rate limiting, flakes, or other reasons. Applies expontential
     backoff with jitter to retry the call."""
 
+    @retry(
+        tries=tries,
+        delay=delay,
+        max_delay=max_delay,
+        backoff=backoff,
+        jitter=jitter,
+        logger=cast(Logger, logger),
+    )
     def retry_with_default(func: F) -> F:
-        @retry(
-            tries=tries,
-            delay=delay,
-            max_delay=max_delay,
-            backoff=backoff,
-            jitter=jitter,
-            logger=cast(Logger, logger),
-        )
         def wrapped_func(*args: list, **kwargs: dict[str, Any]) -> Any:
             return func(*args, **kwargs)
 

backend/danswer/connectors/discourse/connector.py

@@ -14,6 +14,7 @@ from danswer.connectors.cross_connector_utils.miscellaneous_utils import time_st
 from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
     rate_limit_builder,
 )
+from danswer.connectors.cross_connector_utils.retry_wrapper import retry_builder
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import PollConnector
 from danswer.connectors.interfaces import SecondsSinceUnixEpoch
@@ -23,7 +24,6 @@ from danswer.connectors.models import Document
 from danswer.connectors.models import Section
 from danswer.file_processing.html_utils import parse_html_page_basic
 from danswer.utils.logger import setup_logger
-from danswer.utils.retry_wrapper import retry_builder
 
 logger = setup_logger()
 

backend/danswer/connectors/document360/connector.py

@@ -11,6 +11,7 @@ from danswer.configs.constants import DocumentSource
 from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
     rate_limit_builder,
 )
+from danswer.connectors.cross_connector_utils.retry_wrapper import retry_builder
 from danswer.connectors.document360.utils import flatten_child_categories
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import LoadConnector
@@ -21,7 +22,6 @@ from danswer.connectors.models import ConnectorMissingCredentialError
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
 from danswer.file_processing.html_utils import parse_html_page_basic
-from danswer.utils.retry_wrapper import retry_builder
 
 # Limitations and Potential Improvements
 # 1. The "Categories themselves contain potentially relevant information" but they're not pulled in

backend/danswer/connectors/factory.py

@@ -4,7 +4,6 @@ from typing import Type
 from sqlalchemy.orm import Session
 
 from danswer.configs.constants import DocumentSource
-from danswer.configs.constants import DocumentSourceRequiringTenantContext
 from danswer.connectors.asana.connector import AsanaConnector
 from danswer.connectors.axero.connector import AxeroConnector
 from danswer.connectors.blob.connector import BlobStorageConnector
@@ -64,7 +63,7 @@ def identify_connector_class(
         DocumentSource.SLACK: {
             InputType.LOAD_STATE: SlackLoadConnector,
             InputType.POLL: SlackPollConnector,
-            InputType.SLIM_RETRIEVAL: SlackPollConnector,
+            InputType.PRUNE: SlackPollConnector,
         },
         DocumentSource.GITHUB: GithubConnector,
         DocumentSource.GMAIL: GmailConnector,
@@ -135,13 +134,8 @@ def instantiate_connector(
     input_type: InputType,
     connector_specific_config: dict[str, Any],
     credential: Credential,
-    tenant_id: str | None = None,
 ) -> BaseConnector:
     connector_class = identify_connector_class(source, input_type)
-
-    if source in DocumentSourceRequiringTenantContext:
-        connector_specific_config["tenant_id"] = tenant_id
-
     connector = connector_class(**connector_specific_config)
     new_credentials = connector.load_credentials(credential.credential_json)
 

backend/danswer/connectors/file/connector.py

@@ -16,7 +16,7 @@ from danswer.connectors.interfaces import LoadConnector
 from danswer.connectors.models import BasicExpertInfo
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.file_processing.extract_file_text import check_file_ext_is_valid
 from danswer.file_processing.extract_file_text import detect_encoding
 from danswer.file_processing.extract_file_text import extract_file_text
@@ -27,8 +27,6 @@ from danswer.file_processing.extract_file_text import read_pdf_file
 from danswer.file_processing.extract_file_text import read_text_file
 from danswer.file_store.file_store import get_default_file_store
 from danswer.utils.logger import setup_logger
-from shared_configs.configs import CURRENT_TENANT_ID_CONTEXTVAR
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 
 logger = setup_logger()
 
@@ -161,12 +159,10 @@ class LocalFileConnector(LoadConnector):
     def __init__(
         self,
         file_locations: list[Path | str],
-        tenant_id: str = POSTGRES_DEFAULT_SCHEMA,
         batch_size: int = INDEX_BATCH_SIZE,
     ) -> None:
         self.file_locations = [Path(file_location) for file_location in file_locations]
         self.batch_size = batch_size
-        self.tenant_id = tenant_id
         self.pdf_pass: str | None = None
 
     def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
@@ -175,9 +171,7 @@ class LocalFileConnector(LoadConnector):
 
     def load_from_state(self) -> GenerateDocumentsOutput:
         documents: list[Document] = []
-        token = CURRENT_TENANT_ID_CONTEXTVAR.set(self.tenant_id)
-
-        with get_session_with_tenant(self.tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             for file_path in self.file_locations:
                 current_datetime = datetime.now(timezone.utc)
                 files = _read_files_and_metadata(
@@ -199,8 +193,6 @@ class LocalFileConnector(LoadConnector):
             if documents:
                 yield documents
 
-        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
-
 
 if __name__ == "__main__":
     connector = LocalFileConnector(file_locations=[os.environ["TEST_FILE"]])

backend/danswer/connectors/google_drive/connector.py

@@ -19,6 +19,7 @@ from danswer.configs.app_configs import GOOGLE_DRIVE_ONLY_ORG_PUBLIC
 from danswer.configs.app_configs import INDEX_BATCH_SIZE
 from danswer.configs.constants import DocumentSource
 from danswer.configs.constants import IGNORE_FOR_QA
+from danswer.connectors.cross_connector_utils.retry_wrapper import retry_builder
 from danswer.connectors.google_drive.connector_auth import get_google_drive_creds
 from danswer.connectors.google_drive.constants import (
     DB_CREDENTIALS_DICT_DELEGATED_USER_KEY,
@@ -39,7 +40,6 @@ from danswer.file_processing.unstructured import get_unstructured_api_key
 from danswer.file_processing.unstructured import unstructured_to_text
 from danswer.utils.batching import batch_generator
 from danswer.utils.logger import setup_logger
-from danswer.utils.retry_wrapper import retry_builder
 
 logger = setup_logger()
 
@@ -462,34 +462,8 @@ class GoogleDriveConnector(LoadConnector, PollConnector):
                             for permission in file["permissions"]
                         ):
                             continue
-                    try:
-                        text_contents = extract_text(file, service) or ""
-                    except HttpError as e:
-                        reason = (
-                            e.error_details[0]["reason"]
-                            if e.error_details
-                            else e.reason
-                        )
-                        message = (
-                            e.error_details[0]["message"]
-                            if e.error_details
-                            else e.reason
-                        )
 
-                        # these errors don't represent a failure in the connector, but simply files
-                        # that can't / shouldn't be indexed
-                        ERRORS_TO_CONTINUE_ON = [
-                            "cannotExportFile",
-                            "exportSizeLimitExceeded",
-                            "cannotDownloadFile",
-                        ]
-                        if e.status_code == 403 and reason in ERRORS_TO_CONTINUE_ON:
-                            logger.warning(
-                                f"Could not export file '{file['name']}' due to '{message}', skipping..."
-                            )
-                            continue
-
-                        raise
+                    text_contents = extract_text(file, service) or ""
 
                     doc_batch.append(
                         Document(

backend/danswer/connectors/interfaces.py

@@ -3,18 +3,14 @@ from collections.abc import Iterator
 from typing import Any
 
 from danswer.connectors.models import Document
-from danswer.connectors.models import SlimDocument
 
 
 SecondsSinceUnixEpoch = float
 
 GenerateDocumentsOutput = Iterator[list[Document]]
-GenerateSlimDocumentOutput = Iterator[list[SlimDocument]]
 
 
 class BaseConnector(abc.ABC):
-    REDIS_KEY_PREFIX = "da_connector_data:"
-
     @abc.abstractmethod
     def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
         raise NotImplementedError
@@ -54,9 +50,9 @@ class PollConnector(BaseConnector):
         raise NotImplementedError
 
 
-class SlimConnector(BaseConnector):
+class IdConnector(BaseConnector):
     @abc.abstractmethod
-    def retrieve_all_slim_documents(self) -> GenerateSlimDocumentOutput:
+    def retrieve_all_source_ids(self) -> set[str]:
         raise NotImplementedError
 
 

backend/danswer/connectors/loopio/connector.py

@@ -161,7 +161,7 @@ class LoopioConnector(LoadConnector, PollConnector):
                 ]
                 doc_batch.append(
                     Document(
-                        id=str(entry["id"]),
+                        id=entry["id"],
                         sections=[Section(link=link, text=content_text)],
                         source=DocumentSource.LOOPIO,
                         semantic_identifier=questions[0],

backend/danswer/connectors/mediawiki/family.py

@@ -45,7 +45,8 @@ class FamilyFileGeneratorInMemory(generate_family_file.FamilyFileGenerator):
 
         if any(x not in generate_family_file.NAME_CHARACTERS for x in name):
             raise ValueError(
-                f'ERROR: Name of family "{name}" must be ASCII letters and digits [a-zA-Z0-9]',
+                'ERROR: Name of family "{}" must be ASCII letters and digits [a-zA-Z0-9]',
+                name,
             )
 
         if isinstance(dointerwiki, bool):

backend/danswer/connectors/mediawiki/wiki.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import datetime
 import itertools
 from collections.abc import Generator
-from collections.abc import Iterator
 from typing import Any
 from typing import ClassVar
 
@@ -20,9 +19,6 @@ from danswer.connectors.interfaces import SecondsSinceUnixEpoch
 from danswer.connectors.mediawiki.family import family_class_dispatch
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
 
 
 def pywikibot_timestamp_to_utc_datetime(
@@ -78,7 +74,7 @@ def get_doc_from_page(
         sections=sections,
         semantic_identifier=page.title(),
         metadata={"categories": [category.title() for category in page.categories()]},
-        id=f"MEDIAWIKI_{page.pageid}_{page.full_url()}",
+        id=page.pageid,
     )
 
 
@@ -121,18 +117,13 @@ class MediaWikiConnector(LoadConnector, PollConnector):
 
         # short names can only have ascii letters and digits
 
-        self.family = family_class_dispatch(hostname, "WikipediaConnector")()
+        self.family = family_class_dispatch(hostname, "Wikipedia Connector")()
         self.site = pywikibot.Site(fam=self.family, code=language_code)
         self.categories = [
             pywikibot.Category(self.site, f"Category:{category.replace(' ', '_')}")
             for category in categories
         ]
-
-        self.pages = []
-        for page in pages:
-            if not page:
-                continue
-            self.pages.append(pywikibot.Page(self.site, page))
+        self.pages = [pywikibot.Page(self.site, page) for page in pages]
 
     def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
         """Load credentials for a MediaWiki site.
@@ -178,13 +169,8 @@ class MediaWikiConnector(LoadConnector, PollConnector):
         ]
 
         # Since we can specify both individual pages and categories, we need to iterate over all of them.
-        all_pages: Iterator[pywikibot.Page] = itertools.chain(
-            self.pages, *category_pages
-        )
+        all_pages = itertools.chain(self.pages, *category_pages)
         for page in all_pages:
-            logger.info(
-                f"MediaWikiConnector: title='{page.title()}' url={page.full_url()}"
-            )
             doc_batch.append(
                 get_doc_from_page(page, self.site, self.document_source_type)
             )

backend/danswer/connectors/models.py

@@ -14,7 +14,7 @@ class InputType(str, Enum):
     LOAD_STATE = "load_state"  # e.g. loading a current full state or a save state, such as from a file
     POLL = "poll"  # e.g. calling an API to get all documents in the last hour
     EVENT = "event"  # e.g. registered an endpoint as a listener, and processing connector events
-    SLIM_RETRIEVAL = "slim_retrieval"
+    PRUNE = "prune"
 
 
 class ConnectorMissingCredentialError(PermissionError):
@@ -169,11 +169,6 @@ class Document(DocumentBase):
         )
 
 
-class SlimDocument(BaseModel):
-    id: str
-    perm_sync_data: Any | None = None
-
-
 class DocumentErrorSummary(BaseModel):
     id: str
     semantic_id: str

backend/danswer/connectors/notion/connector.py

@@ -29,9 +29,6 @@ logger = setup_logger()
 _NOTION_CALL_TIMEOUT = 30  # 30 seconds
 
 
-# TODO: Tables need to be ingested, Pages need to have their metadata ingested
-
-
 @dataclass
 class NotionPage:
     """Represents a Notion Page object"""
@@ -43,8 +40,6 @@ class NotionPage:
     properties: dict[str, Any]
     url: str
 
-    database_name: str | None  # Only applicable to the database type page (wiki)
-
     def __init__(self, **kwargs: dict[str, Any]) -> None:
         names = set([f.name for f in fields(self)])
         for k, v in kwargs.items():
@@ -52,17 +47,6 @@ class NotionPage:
                 setattr(self, k, v)
 
 
-@dataclass
-class NotionBlock:
-    """Represents a Notion Block object"""
-
-    id: str  # Used for the URL
-    text: str
-    # In a plaintext representation of the page, how this block should be joined
-    # with the existing text up to this point, separated out from text for clarity
-    prefix: str
-
-
 @dataclass
 class NotionSearchResponse:
     """Represents the response from the Notion Search API"""
@@ -78,6 +62,7 @@ class NotionSearchResponse:
                 setattr(self, k, v)
 
 
+# TODO - Add the ability to optionally limit to specific Notion databases
 class NotionConnector(LoadConnector, PollConnector):
     """Notion Page connector that reads all Notion pages
     this integration has been granted access to.
@@ -141,46 +126,20 @@ class NotionConnector(LoadConnector, PollConnector):
 
     @retry(tries=3, delay=1, backoff=2)
     def _fetch_page(self, page_id: str) -> NotionPage:
-        """Fetch a page from its ID via the Notion API, retry with database if page fetch fails."""
+        """Fetch a page from it's ID via the Notion API."""
         logger.debug(f"Fetching page for ID '{page_id}'")
-        page_url = f"https://api.notion.com/v1/pages/{page_id}"
+        block_url = f"https://api.notion.com/v1/pages/{page_id}"
         res = rl_requests.get(
-            page_url,
+            block_url,
             headers=self.headers,
             timeout=_NOTION_CALL_TIMEOUT,
         )
         try:
             res.raise_for_status()
         except Exception as e:
-            logger.warning(
-                f"Failed to fetch page, trying database for ID '{page_id}'. Exception: {e}"
-            )
-            # Try fetching as a database if page fetch fails, this happens if the page is set to a wiki
-            # it becomes a database from the notion perspective
-            return self._fetch_database_as_page(page_id)
-        return NotionPage(**res.json())
-
-    @retry(tries=3, delay=1, backoff=2)
-    def _fetch_database_as_page(self, database_id: str) -> NotionPage:
-        """Attempt to fetch a database as a page."""
-        logger.debug(f"Fetching database for ID '{database_id}' as a page")
-        database_url = f"https://api.notion.com/v1/databases/{database_id}"
-        res = rl_requests.get(
-            database_url,
-            headers=self.headers,
-            timeout=_NOTION_CALL_TIMEOUT,
-        )
-        try:
-            res.raise_for_status()
-        except Exception as e:
-            logger.exception(f"Error fetching database as page - {res.json()}")
+            logger.exception(f"Error fetching page - {res.json()}")
             raise e
-        database_name = res.json().get("title")
-        database_name = (
-            database_name[0].get("text", {}).get("content") if database_name else None
-        )
-
-        return NotionPage(**res.json(), database_name=database_name)
+        return NotionPage(**res.json())
 
     @retry(tries=3, delay=1, backoff=2)
     def _fetch_database(
@@ -212,86 +171,8 @@ class NotionConnector(LoadConnector, PollConnector):
             raise e
         return res.json()
 
-    @staticmethod
-    def _properties_to_str(properties: dict[str, Any]) -> str:
-        """Converts Notion properties to a string"""
-
-        def _recurse_properties(inner_dict: dict[str, Any]) -> str | None:
-            while "type" in inner_dict:
-                type_name = inner_dict["type"]
-                inner_dict = inner_dict[type_name]
-
-                # If the innermost layer is None, the value is not set
-                if not inner_dict:
-                    return None
-
-                if isinstance(inner_dict, list):
-                    list_properties = [
-                        _recurse_properties(item) for item in inner_dict if item
-                    ]
-                    return (
-                        ", ".join(
-                            [
-                                list_property
-                                for list_property in list_properties
-                                if list_property
-                            ]
-                        )
-                        or None
-                    )
-
-            # TODO there may be more types to handle here
-            if isinstance(inner_dict, str):
-                # For some objects the innermost value could just be a string, not sure what causes this
-                return inner_dict
-
-            elif isinstance(inner_dict, dict):
-                if "name" in inner_dict:
-                    return inner_dict["name"]
-                if "content" in inner_dict:
-                    return inner_dict["content"]
-                start = inner_dict.get("start")
-                end = inner_dict.get("end")
-                if start is not None:
-                    if end is not None:
-                        return f"{start} - {end}"
-                    return start
-                elif end is not None:
-                    return f"Until {end}"
-
-                if "id" in inner_dict:
-                    # This is not useful to index, it's a reference to another Notion object
-                    # and this ID value in plaintext is useless outside of the Notion context
-                    logger.debug("Skipping Notion object id field property")
-                    return None
-
-            logger.debug(f"Unreadable property from innermost prop: {inner_dict}")
-            return None
-
-        result = ""
-        for prop_name, prop in properties.items():
-            if not prop:
-                continue
-
-            try:
-                inner_value = _recurse_properties(prop)
-            except Exception as e:
-                # This is not a critical failure, these properties are not the actual contents of the page
-                # more similar to metadata
-                logger.warning(f"Error recursing properties for {prop_name}: {e}")
-                continue
-            # Not a perfect way to format Notion database tables but there's no perfect representation
-            # since this must be represented as plaintext
-            if inner_value:
-                result += f"{prop_name}: {inner_value}\t"
-
-        return result
-
-    def _read_pages_from_database(
-        self, database_id: str
-    ) -> tuple[list[NotionBlock], list[str]]:
-        """Returns a list of top level blocks and all page IDs in the database"""
-        result_blocks: list[NotionBlock] = []
+    def _read_pages_from_database(self, database_id: str) -> list[str]:
+        """Returns a list of all page IDs in the database"""
         result_pages: list[str] = []
         cursor = None
         while True:
@@ -300,34 +181,29 @@ class NotionConnector(LoadConnector, PollConnector):
             for result in data["results"]:
                 obj_id = result["id"]
                 obj_type = result["object"]
-                text = self._properties_to_str(result.get("properties", {}))
-                if text:
-                    result_blocks.append(NotionBlock(id=obj_id, text=text, prefix="\n"))
-
-                if self.recursive_index_enabled:
-                    if obj_type == "page":
-                        logger.debug(
-                            f"Found page with ID '{obj_id}' in database '{database_id}'"
-                        )
-                        result_pages.append(result["id"])
-                    elif obj_type == "database":
-                        logger.debug(
-                            f"Found database with ID '{obj_id}' in database '{database_id}'"
-                        )
-                        # The inner contents are ignored at this level
-                        _, child_pages = self._read_pages_from_database(obj_id)
-                        result_pages.extend(child_pages)
+                if obj_type == "page":
+                    logger.debug(
+                        f"Found page with ID '{obj_id}' in database '{database_id}'"
+                    )
+                    result_pages.append(result["id"])
+                elif obj_type == "database":
+                    logger.debug(
+                        f"Found database with ID '{obj_id}' in database '{database_id}'"
+                    )
+                    result_pages.extend(self._read_pages_from_database(obj_id))
 
             if data["next_cursor"] is None:
                 break
 
             cursor = data["next_cursor"]
 
-        return result_blocks, result_pages
+        return result_pages
 
-    def _read_blocks(self, base_block_id: str) -> tuple[list[NotionBlock], list[str]]:
-        """Reads all child blocks for the specified block, returns a list of blocks and child page ids"""
-        result_blocks: list[NotionBlock] = []
+    def _read_blocks(
+        self, base_block_id: str
+    ) -> tuple[list[tuple[str, str]], list[str]]:
+        """Reads all child blocks for the specified block"""
+        result_lines: list[tuple[str, str]] = []
         child_pages: list[str] = []
         cursor = None
         while True:
@@ -335,7 +211,7 @@ class NotionConnector(LoadConnector, PollConnector):
 
             # this happens when a block is not shared with the integration
             if data is None:
-                return result_blocks, child_pages
+                return result_lines, child_pages
 
             for result in data["results"]:
                 logger.debug(
@@ -379,70 +255,46 @@ class NotionConnector(LoadConnector, PollConnector):
 
                 if result["has_children"]:
                     if result_type == "child_page":
-                        # Child pages will not be included at this top level, it will be a separate document
                         child_pages.append(result_block_id)
                     else:
                         logger.debug(f"Entering sub-block: {result_block_id}")
-                        subblocks, subblock_child_pages = self._read_blocks(
+                        subblock_result_lines, subblock_child_pages = self._read_blocks(
                             result_block_id
                         )
                         logger.debug(f"Finished sub-block: {result_block_id}")
-                        result_blocks.extend(subblocks)
+                        result_lines.extend(subblock_result_lines)
                         child_pages.extend(subblock_child_pages)
 
-                if result_type == "child_database":
-                    inner_blocks, inner_child_pages = self._read_pages_from_database(
-                        result_block_id
-                    )
-                    # A database on a page often looks like a table, we need to include it for the contents
-                    # of the page but the children (cells) should be processed as other Documents
-                    result_blocks.extend(inner_blocks)
+                if result_type == "child_database" and self.recursive_index_enabled:
+                    child_pages.extend(self._read_pages_from_database(result_block_id))
 
-                    if self.recursive_index_enabled:
-                        child_pages.extend(inner_child_pages)
-
-                if cur_result_text_arr:
-                    new_block = NotionBlock(
-                        id=result_block_id,
-                        text="\n".join(cur_result_text_arr),
-                        prefix="\n",
-                    )
-                    result_blocks.append(new_block)
+                cur_result_text = "\n".join(cur_result_text_arr)
+                if cur_result_text:
+                    result_lines.append((cur_result_text, result_block_id))
 
             if data["next_cursor"] is None:
                 break
 
             cursor = data["next_cursor"]
 
-        return result_blocks, child_pages
+        return result_lines, child_pages
 
-    def _read_page_title(self, page: NotionPage) -> str | None:
+    def _read_page_title(self, page: NotionPage) -> str:
         """Extracts the title from a Notion page"""
         page_title = None
-        if hasattr(page, "database_name") and page.database_name:
-            return page.database_name
         for _, prop in page.properties.items():
             if prop["type"] == "title" and len(prop["title"]) > 0:
                 page_title = " ".join([t["plain_text"] for t in prop["title"]]).strip()
                 break
-
+        if page_title is None:
+            page_title = f"Untitled Page [{page.id}]"
         return page_title
 
     def _read_pages(
         self,
         pages: list[NotionPage],
     ) -> Generator[Document, None, None]:
-        """Reads pages for rich text content and generates Documents
-
-        Note that a page which is turned into a "wiki" becomes a database but both top level pages and top level databases
-        do not seem to have any properties associated with them.
-
-        Pages that are part of a database can have properties which are like the values of the row in the "database" table
-        in which they exist
-
-        This is not clearly outlined in the Notion API docs but it is observable empirically.
-        https://developers.notion.com/docs/working-with-page-content
-        """
+        """Reads pages for rich text content and generates Documents"""
         all_child_page_ids: list[str] = []
         for page in pages:
             if page.id in self.indexed_pages:
@@ -452,23 +304,18 @@ class NotionConnector(LoadConnector, PollConnector):
             logger.info(f"Reading page with ID '{page.id}', with url {page.url}")
             page_blocks, child_page_ids = self._read_blocks(page.id)
             all_child_page_ids.extend(child_page_ids)
-
-            if not page_blocks:
-                continue
-
-            page_title = (
-                self._read_page_title(page) or f"Untitled Page with ID {page.id}"
-            )
-
+            page_title = self._read_page_title(page)
             yield (
                 Document(
                     id=page.id,
-                    sections=[
+                    # Will add title to the first section later in processing
+                    sections=[Section(link=page.url, text="")]
+                    + [
                         Section(
-                            link=f"{page.url}#{block.id.replace('-', '')}",
-                            text=block.prefix + block.text,
+                            link=f"{page.url}#{block_id.replace('-', '')}",
+                            text=block_text,
                         )
-                        for block in page_blocks
+                        for block_text, block_id in page_blocks
                     ],
                     source=DocumentSource.NOTION,
                     semantic_identifier=page_title,

backend/danswer/connectors/salesforce/connector.py

@@ -11,25 +11,17 @@ from danswer.configs.app_configs import INDEX_BATCH_SIZE
 from danswer.configs.constants import DocumentSource
 from danswer.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc
 from danswer.connectors.interfaces import GenerateDocumentsOutput
-from danswer.connectors.interfaces import GenerateSlimDocumentOutput
+from danswer.connectors.interfaces import IdConnector
 from danswer.connectors.interfaces import LoadConnector
 from danswer.connectors.interfaces import PollConnector
 from danswer.connectors.interfaces import SecondsSinceUnixEpoch
-from danswer.connectors.interfaces import SlimConnector
 from danswer.connectors.models import BasicExpertInfo
 from danswer.connectors.models import ConnectorMissingCredentialError
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
-from danswer.connectors.models import SlimDocument
 from danswer.connectors.salesforce.utils import extract_dict_text
 from danswer.utils.logger import setup_logger
 
-
-# TODO: this connector does not work well at large scales
-# the large query against a large Salesforce instance has been reported to take 1.5 hours.
-# Additionally it seems to eat up more memory over time if the connection is long running (again a scale issue).
-
-
 DEFAULT_PARENT_OBJECT_TYPES = ["Account"]
 MAX_QUERY_LENGTH = 10000  # max query length is 20,000 characters
 ID_PREFIX = "SALESFORCE_"
@@ -37,7 +29,7 @@ ID_PREFIX = "SALESFORCE_"
 logger = setup_logger()
 
 
-class SalesforceConnector(LoadConnector, PollConnector, SlimConnector):
+class SalesforceConnector(LoadConnector, PollConnector, IdConnector):
     def __init__(
         self,
         batch_size: int = INDEX_BATCH_SIZE,
@@ -251,22 +243,19 @@ class SalesforceConnector(LoadConnector, PollConnector, SlimConnector):
         end_datetime = datetime.utcfromtimestamp(end)
         return self._fetch_from_salesforce(start=start_datetime, end=end_datetime)
 
-    def retrieve_all_slim_documents(self) -> GenerateSlimDocumentOutput:
+    def retrieve_all_source_ids(self) -> set[str]:
         if self.sf_client is None:
             raise ConnectorMissingCredentialError("Salesforce")
-        doc_metadata_list: list[SlimDocument] = []
+        all_retrieved_ids: set[str] = set()
         for parent_object_type in self.parent_object_list:
             query = f"SELECT Id FROM {parent_object_type}"
             query_result = self.sf_client.query_all(query)
-            doc_metadata_list.extend(
-                SlimDocument(
-                    id=f"{ID_PREFIX}{instance_dict.get('Id', '')}",
-                    perm_sync_data={},
-                )
+            all_retrieved_ids.update(
+                f"{ID_PREFIX}{instance_dict.get('Id', '')}"
                 for instance_dict in query_result["records"]
             )
 
-        yield doc_metadata_list
+        return all_retrieved_ids
 
 
 if __name__ == "__main__":

backend/danswer/connectors/slack/connector.py

@@ -13,15 +13,13 @@ from danswer.configs.app_configs import ENABLE_EXPENSIVE_EXPERT_CALLS
 from danswer.configs.app_configs import INDEX_BATCH_SIZE
 from danswer.configs.constants import DocumentSource
 from danswer.connectors.interfaces import GenerateDocumentsOutput
-from danswer.connectors.interfaces import GenerateSlimDocumentOutput
+from danswer.connectors.interfaces import IdConnector
 from danswer.connectors.interfaces import PollConnector
 from danswer.connectors.interfaces import SecondsSinceUnixEpoch
-from danswer.connectors.interfaces import SlimConnector
 from danswer.connectors.models import BasicExpertInfo
 from danswer.connectors.models import ConnectorMissingCredentialError
 from danswer.connectors.models import Document
 from danswer.connectors.models import Section
-from danswer.connectors.models import SlimDocument
 from danswer.connectors.slack.utils import expert_info_from_slack_id
 from danswer.connectors.slack.utils import get_message_link
 from danswer.connectors.slack.utils import make_paginated_slack_api_call_w_retries
@@ -328,7 +326,7 @@ def _get_all_doc_ids(
     channels: list[str] | None = None,
     channel_name_regex_enabled: bool = False,
     msg_filter_func: Callable[[MessageType], bool] = default_msg_filter,
-) -> GenerateSlimDocumentOutput:
+) -> set[str]:
     """
     Get all document ids in the workspace, channel by channel
     This is pretty identical to get_all_docs, but it returns a set of ids instead of documents
@@ -340,14 +338,13 @@ def _get_all_doc_ids(
         all_channels, channels, channel_name_regex_enabled
     )
 
+    all_doc_ids = set()
     for channel in filtered_channels:
-        channel_id = channel["id"]
         channel_message_batches = get_channel_messages(
             client=client,
             channel=channel,
         )
 
-        message_ts_set: set[str] = set()
         for message_batch in channel_message_batches:
             for message in message_batch:
                 if msg_filter_func(message):
@@ -356,21 +353,12 @@ def _get_all_doc_ids(
                 # The document id is the channel id and the ts of the first message in the thread
                 # Since we already have the first message of the thread, we dont have to
                 # fetch the thread for id retrieval, saving time and API calls
-                message_ts_set.add(message["ts"])
+                all_doc_ids.add(f"{channel['id']}__{message['ts']}")
 
-        channel_metadata_list: list[SlimDocument] = []
-        for message_ts in message_ts_set:
-            channel_metadata_list.append(
-                SlimDocument(
-                    id=f"{channel_id}__{message_ts}",
-                    perm_sync_data={"channel_id": channel_id},
-                )
-            )
-
-        yield channel_metadata_list
+    return all_doc_ids
 
 
-class SlackPollConnector(PollConnector, SlimConnector):
+class SlackPollConnector(PollConnector, IdConnector):
     def __init__(
         self,
         workspace: str,
@@ -391,7 +379,7 @@ class SlackPollConnector(PollConnector, SlimConnector):
         self.client = WebClient(token=bot_token)
         return None
 
-    def retrieve_all_slim_documents(self) -> GenerateSlimDocumentOutput:
+    def retrieve_all_source_ids(self) -> set[str]:
         if self.client is None:
             raise ConnectorMissingCredentialError("Slack")
 

backend/danswer/connectors/slack/utils.py

@@ -10,9 +10,9 @@ from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
 from slack_sdk.web import SlackResponse
 
+from danswer.connectors.cross_connector_utils.retry_wrapper import retry_builder
 from danswer.connectors.models import BasicExpertInfo
 from danswer.utils.logger import setup_logger
-from danswer.utils.retry_wrapper import retry_builder
 
 logger = setup_logger()
 

backend/danswer/connectors/web/connector.py

@@ -128,9 +128,6 @@ def get_internal_links(
         if not href:
             continue
 
-        # Account for malformed backslashes in URLs
-        href = href.replace("\\", "/")
-
         if should_ignore_pound and "#" in href:
             href = href.split("#")[0]
 
@@ -373,7 +370,7 @@ class WebConnector(LoadConnector):
                 page.close()
             except Exception as e:
                 last_error = f"Failed to fetch '{current_url}': {e}"
-                logger.exception(last_error)
+                logger.error(last_error)
                 playwright.stop()
                 restart_playwright = True
                 continue

backend/danswer/danswerbot/slack/formatting.py

@@ -1,66 +0,0 @@
-from mistune import Markdown  # type: ignore
-from mistune import Renderer  # type: ignore
-
-
-def format_slack_message(message: str | None) -> str:
-    renderer = Markdown(renderer=SlackRenderer())
-    return renderer.render(message)
-
-
-class SlackRenderer(Renderer):
-    SPECIALS: dict[str, str] = {"&": "&amp;", "<": "&lt;", ">": "&gt;"}
-
-    def escape_special(self, text: str) -> str:
-        for special, replacement in self.SPECIALS.items():
-            text = text.replace(special, replacement)
-        return text
-
-    def header(self, text: str, level: int, raw: str | None = None) -> str:
-        return f"*{text}*\n"
-
-    def emphasis(self, text: str) -> str:
-        return f"_{text}_"
-
-    def double_emphasis(self, text: str) -> str:
-        return f"*{text}*"
-
-    def strikethrough(self, text: str) -> str:
-        return f"~{text}~"
-
-    def list(self, body: str, ordered: bool = True) -> str:
-        lines = body.split("\n")
-        count = 0
-        for i, line in enumerate(lines):
-            if line.startswith("li: "):
-                count += 1
-                prefix = f"{count}. " if ordered else "• "
-                lines[i] = f"{prefix}{line[4:]}"
-        return "\n".join(lines)
-
-    def list_item(self, text: str) -> str:
-        return f"li: {text}\n"
-
-    def link(self, link: str, title: str | None, content: str | None) -> str:
-        escaped_link = self.escape_special(link)
-        if content:
-            return f"<{escaped_link}|{content}>"
-        if title:
-            return f"<{escaped_link}|{title}>"
-        return f"<{escaped_link}>"
-
-    def image(self, src: str, title: str | None, text: str | None) -> str:
-        escaped_src = self.escape_special(src)
-        display_text = title or text
-        return f"<{escaped_src}|{display_text}>" if display_text else f"<{escaped_src}>"
-
-    def codespan(self, text: str) -> str:
-        return f"`{text}`"
-
-    def block_code(self, text: str, lang: str | None) -> str:
-        return f"```\n{text}\n```\n"
-
-    def paragraph(self, text: str) -> str:
-        return f"{text}\n"
-
-    def autolink(self, link: str, is_email: bool) -> str:
-        return link if is_email else self.link(link, None, None)

backend/danswer/danswerbot/slack/handlers/handle_buttons.py

@@ -4,7 +4,9 @@ from typing import cast
 from slack_sdk import WebClient
 from slack_sdk.models.blocks import SectionBlock
 from slack_sdk.models.views import View
+from slack_sdk.socket_mode import SocketModeClient
 from slack_sdk.socket_mode.request import SocketModeRequest
+from sqlalchemy.orm import Session
 
 from danswer.configs.constants import MessageType
 from danswer.configs.constants import SearchFeedbackType
@@ -33,22 +35,20 @@ from danswer.danswerbot.slack.utils import get_channel_name_from_id
 from danswer.danswerbot.slack.utils import get_feedback_visibility
 from danswer.danswerbot.slack.utils import read_slack_thread
 from danswer.danswerbot.slack.utils import respond_in_thread
-from danswer.danswerbot.slack.utils import TenantSocketModeClient
 from danswer.danswerbot.slack.utils import update_emote_react
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.feedback import create_chat_message_feedback
 from danswer.db.feedback import create_doc_retrieval_feedback
 from danswer.document_index.document_index_utils import get_both_index_names
 from danswer.document_index.factory import get_default_document_index
 from danswer.utils.logger import setup_logger
 
-
 logger = setup_logger()
 
 
 def handle_doc_feedback_button(
     req: SocketModeRequest,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
 ) -> None:
     if not (actions := req.payload.get("actions")):
         logger.error("Missing actions. Unable to build the source feedback view")
@@ -81,7 +81,7 @@ def handle_doc_feedback_button(
 
 def handle_generate_answer_button(
     req: SocketModeRequest,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
 ) -> None:
     channel_id = req.payload["channel"]["id"]
     channel_name = req.payload["channel"]["name"]
@@ -116,7 +116,7 @@ def handle_generate_answer_button(
         thread_ts=thread_ts,
     )
 
-    with get_session_with_tenant(client.tenant_id) as db_session:
+    with Session(get_sqlalchemy_engine()) as db_session:
         slack_bot_config = get_slack_bot_config_for_channel(
             channel_name=channel_name, db_session=db_session
         )
@@ -136,7 +136,6 @@ def handle_generate_answer_button(
             slack_bot_config=slack_bot_config,
             receiver_ids=None,
             client=client.web_client,
-            tenant_id=client.tenant_id,
             channel=channel_id,
             logger=logger,
             feedback_reminder_id=None,
@@ -151,11 +150,12 @@ def handle_slack_feedback(
     user_id_to_post_confirmation: str,
     channel_id_to_post_confirmation: str,
     thread_ts_to_post_confirmation: str,
-    tenant_id: str | None,
 ) -> None:
+    engine = get_sqlalchemy_engine()
+
     message_id, doc_id, doc_rank = decompose_action_id(feedback_id)
 
-    with get_session_with_tenant(tenant_id) as db_session:
+    with Session(engine) as db_session:
         if feedback_type in [LIKE_BLOCK_ACTION_ID, DISLIKE_BLOCK_ACTION_ID]:
             create_chat_message_feedback(
                 is_positive=feedback_type == LIKE_BLOCK_ACTION_ID,
@@ -232,7 +232,7 @@ def handle_slack_feedback(
 
 def handle_followup_button(
     req: SocketModeRequest,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
 ) -> None:
     action_id = None
     if actions := req.payload.get("actions"):
@@ -252,7 +252,7 @@ def handle_followup_button(
 
     tag_ids: list[str] = []
     group_ids: list[str] = []
-    with get_session_with_tenant(client.tenant_id) as db_session:
+    with Session(get_sqlalchemy_engine()) as db_session:
         channel_name, is_dm = get_channel_name_from_id(
             client=client.web_client, channel_id=channel_id
         )
@@ -295,7 +295,7 @@ def handle_followup_button(
 
 def get_clicker_name(
     req: SocketModeRequest,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
 ) -> str:
     clicker_name = req.payload.get("user", {}).get("name", "Someone")
     clicker_real_name = None
@@ -316,7 +316,7 @@ def get_clicker_name(
 
 def handle_followup_resolved_button(
     req: SocketModeRequest,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
     immediate: bool = False,
 ) -> None:
     channel_id = req.payload["container"]["channel_id"]

backend/danswer/danswerbot/slack/handlers/handle_message.py

@@ -2,6 +2,7 @@ import datetime
 
 from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
+from sqlalchemy.orm import Session
 
 from danswer.configs.danswerbot_configs import DANSWER_BOT_FEEDBACK_REMINDER
 from danswer.configs.danswerbot_configs import DANSWER_REACT_EMOJI
@@ -18,7 +19,7 @@ from danswer.danswerbot.slack.utils import fetch_user_ids_from_groups
 from danswer.danswerbot.slack.utils import respond_in_thread
 from danswer.danswerbot.slack.utils import slack_usage_report
 from danswer.danswerbot.slack.utils import update_emote_react
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.models import SlackBotConfig
 from danswer.db.users import add_non_web_user_if_not_exists
 from danswer.utils.logger import setup_logger
@@ -109,7 +110,6 @@ def handle_message(
     slack_bot_config: SlackBotConfig | None,
     client: WebClient,
     feedback_reminder_id: str | None,
-    tenant_id: str | None,
 ) -> bool:
     """Potentially respond to the user message depending on filters and if an answer was generated
 
@@ -135,9 +135,7 @@ def handle_message(
         action = "slack_tag_message"
     elif is_bot_dm:
         action = "slack_dm_message"
-    slack_usage_report(
-        action=action, sender_id=sender_id, client=client, tenant_id=tenant_id
-    )
+    slack_usage_report(action=action, sender_id=sender_id, client=client)
 
     document_set_names: list[str] | None = None
     persona = slack_bot_config.persona if slack_bot_config else None
@@ -211,7 +209,7 @@ def handle_message(
     except SlackApiError as e:
         logger.error(f"Was not able to react to user message due to: {e}")
 
-    with get_session_with_tenant(tenant_id) as db_session:
+    with Session(get_sqlalchemy_engine()) as db_session:
         if message_info.email:
             add_non_web_user_if_not_exists(db_session, message_info.email)
 
@@ -237,6 +235,5 @@ def handle_message(
             channel=channel,
             logger=logger,
             feedback_reminder_id=feedback_reminder_id,
-            tenant_id=tenant_id,
         )
         return issue_with_regular_answer

backend/danswer/danswerbot/slack/handlers/handle_regular_answer.py

@@ -5,10 +5,12 @@ from typing import cast
 from typing import Optional
 from typing import TypeVar
 
+from fastapi import HTTPException
 from retry import retry
 from slack_sdk import WebClient
 from slack_sdk.models.blocks import DividerBlock
 from slack_sdk.models.blocks import SectionBlock
+from sqlalchemy.orm import Session
 
 from danswer.configs.app_configs import DISABLE_GENERATIVE_AI
 from danswer.configs.danswerbot_configs import DANSWER_BOT_ANSWER_GENERATION_TIMEOUT
@@ -26,13 +28,12 @@ from danswer.danswerbot.slack.blocks import build_follow_up_block
 from danswer.danswerbot.slack.blocks import build_qa_response_blocks
 from danswer.danswerbot.slack.blocks import build_sources_blocks
 from danswer.danswerbot.slack.blocks import get_restate_blocks
-from danswer.danswerbot.slack.formatting import format_slack_message
 from danswer.danswerbot.slack.handlers.utils import send_team_member_message
 from danswer.danswerbot.slack.models import SlackMessageInfo
 from danswer.danswerbot.slack.utils import respond_in_thread
 from danswer.danswerbot.slack.utils import SlackRateLimiter
 from danswer.danswerbot.slack.utils import update_emote_react
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.models import Persona
 from danswer.db.models import SlackBotConfig
 from danswer.db.models import SlackBotResponseType
@@ -87,7 +88,6 @@ def handle_regular_answer(
     channel: str,
     logger: DanswerLoggingAdapter,
     feedback_reminder_id: str | None,
-    tenant_id: str | None,
     num_retries: int = DANSWER_BOT_NUM_RETRIES,
     answer_generation_timeout: int = DANSWER_BOT_ANSWER_GENERATION_TIMEOUT,
     thread_context_percent: float = DANSWER_BOT_TARGET_CHUNK_PERCENTAGE,
@@ -104,7 +104,8 @@ def handle_regular_answer(
     user = None
     if message_info.is_bot_dm:
         if message_info.email:
-            with get_session_with_tenant(tenant_id) as db_session:
+            engine = get_sqlalchemy_engine()
+            with Session(engine) as db_session:
                 user = get_user_by_email(message_info.email, db_session)
 
     document_set_names: list[str] | None = None
@@ -151,10 +152,14 @@ def handle_regular_answer(
         max_document_tokens: int | None = None
         max_history_tokens: int | None = None
 
-        with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             if len(new_message_request.messages) > 1:
                 if new_message_request.persona_config:
-                    raise RuntimeError("Slack bot does not support persona config")
+                    raise HTTPException(
+                        status_code=403,
+                        detail="Slack bot does not support persona config",
+                    )
+
                 elif new_message_request.persona_id is not None:
                     persona = cast(
                         Persona,
@@ -165,10 +170,6 @@ def handle_regular_answer(
                             get_editable=False,
                         ),
                     )
-                else:
-                    raise RuntimeError(
-                        "No persona id provided, this should never happen."
-                    )
 
                 llm, _ = get_llms_for_persona(persona)
 
@@ -211,7 +212,6 @@ def handle_regular_answer(
                 use_citations=use_citations,
                 danswerbot_flow=True,
             )
-
             if not answer.error_msg:
                 return answer
             else:
@@ -246,7 +246,7 @@ def handle_regular_answer(
         )
 
         # Always apply reranking settings if it exists, this is the non-streaming flow
-        with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             saved_search_settings = get_current_search_settings(db_session)
 
         # This includes throwing out answer via reflexion
@@ -413,11 +413,10 @@ def handle_regular_answer(
 
     # If called with the DanswerBot slash command, the question is lost so we have to reshow it
     restate_question_block = get_restate_blocks(messages[-1].message, is_bot_msg)
-    formatted_answer = format_slack_message(answer.answer) if answer.answer else None
 
     answer_blocks = build_qa_response_blocks(
         message_id=answer.chat_message_id,
-        answer=formatted_answer,
+        answer=answer.answer,
         quotes=answer.quotes.quotes if answer.quotes else None,
         source_filters=retrieval_info.applied_source_filters,
         time_cutoff=retrieval_info.applied_time_cutoff,

backend/danswer/danswerbot/slack/listener.py

@@ -4,8 +4,10 @@ from typing import Any
 from typing import cast
 
 from slack_sdk import WebClient
+from slack_sdk.socket_mode import SocketModeClient
 from slack_sdk.socket_mode.request import SocketModeRequest
 from slack_sdk.socket_mode.response import SocketModeResponse
+from sqlalchemy.orm import Session
 
 from danswer.configs.constants import MessageType
 from danswer.configs.danswerbot_configs import DANSWER_BOT_REPHRASE_MESSAGE
@@ -45,9 +47,7 @@ from danswer.danswerbot.slack.utils import read_slack_thread
 from danswer.danswerbot.slack.utils import remove_danswer_bot_tag
 from danswer.danswerbot.slack.utils import rephrase_slack_message
 from danswer.danswerbot.slack.utils import respond_in_thread
-from danswer.danswerbot.slack.utils import TenantSocketModeClient
-from danswer.db.engine import get_all_tenant_ids
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.search_settings import get_current_search_settings
 from danswer.key_value_store.interface import KvKeyNotFoundError
 from danswer.natural_language_processing.search_nlp_models import EmbeddingModel
@@ -57,7 +57,6 @@ from danswer.search.retrieval.search_runner import download_nltk_data
 from danswer.server.manage.models import SlackBotTokens
 from danswer.utils.logger import setup_logger
 from danswer.utils.variable_functionality import set_is_ee_based_on_env_variable
-from shared_configs.configs import CURRENT_TENANT_ID_CONTEXTVAR
 from shared_configs.configs import MODEL_SERVER_HOST
 from shared_configs.configs import MODEL_SERVER_PORT
 from shared_configs.configs import SLACK_CHANNEL_ID
@@ -81,7 +80,7 @@ _SLACK_GREETINGS_TO_IGNORE = {
 _OFFICIAL_SLACKBOT_USER_ID = "USLACKBOT"
 
 
-def prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -> bool:
+def prefilter_requests(req: SocketModeRequest, client: SocketModeClient) -> bool:
     """True to keep going, False to ignore this Slack request"""
     if req.type == "events_api":
         # Verify channel is valid
@@ -154,7 +153,8 @@ def prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -
                 client=client.web_client, channel_id=channel
             )
 
-            with get_session_with_tenant(client.tenant_id) as db_session:
+            engine = get_sqlalchemy_engine()
+            with Session(engine) as db_session:
                 slack_bot_config = get_slack_bot_config_for_channel(
                     channel_name=channel_name, db_session=db_session
                 )
@@ -221,7 +221,7 @@ def prefilter_requests(req: SocketModeRequest, client: TenantSocketModeClient) -
     return True
 
 
-def process_feedback(req: SocketModeRequest, client: TenantSocketModeClient) -> None:
+def process_feedback(req: SocketModeRequest, client: SocketModeClient) -> None:
     if actions := req.payload.get("actions"):
         action = cast(dict[str, Any], actions[0])
         feedback_type = cast(str, action.get("action_id"))
@@ -243,7 +243,6 @@ def process_feedback(req: SocketModeRequest, client: TenantSocketModeClient) ->
         user_id_to_post_confirmation=user_id,
         channel_id_to_post_confirmation=channel_id,
         thread_ts_to_post_confirmation=thread_ts,
-        tenant_id=client.tenant_id,
     )
 
     query_event_id, _, _ = decompose_action_id(feedback_id)
@@ -251,7 +250,7 @@ def process_feedback(req: SocketModeRequest, client: TenantSocketModeClient) ->
 
 
 def build_request_details(
-    req: SocketModeRequest, client: TenantSocketModeClient
+    req: SocketModeRequest, client: SocketModeClient
 ) -> SlackMessageInfo:
     if req.type == "events_api":
         event = cast(dict[str, Any], req.payload["event"])
@@ -330,7 +329,7 @@ def build_request_details(
 
 def apologize_for_fail(
     details: SlackMessageInfo,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
 ) -> None:
     respond_in_thread(
         client=client.web_client,
@@ -342,13 +341,11 @@ def apologize_for_fail(
 
 def process_message(
     req: SocketModeRequest,
-    client: TenantSocketModeClient,
+    client: SocketModeClient,
     respond_every_channel: bool = DANSWER_BOT_RESPOND_EVERY_CHANNEL,
     notify_no_answer: bool = NOTIFY_SLACKBOT_NO_ANSWER,
 ) -> None:
-    logger.debug(
-        f"Received Slack request of type: '{req.type}' for tenant, {client.tenant_id}"
-    )
+    logger.debug(f"Received Slack request of type: '{req.type}'")
 
     # Throw out requests that can't or shouldn't be handled
     if not prefilter_requests(req, client):
@@ -360,67 +357,59 @@ def process_message(
         client=client.web_client, channel_id=channel
     )
 
-    # Set the current tenant ID at the beginning for all DB calls within this thread
-    if client.tenant_id:
-        logger.info(f"Setting tenant ID to {client.tenant_id}")
-        token = CURRENT_TENANT_ID_CONTEXTVAR.set(client.tenant_id)
-    try:
-        with get_session_with_tenant(client.tenant_id) as db_session:
-            slack_bot_config = get_slack_bot_config_for_channel(
-                channel_name=channel_name, db_session=db_session
-            )
+    engine = get_sqlalchemy_engine()
+    with Session(engine) as db_session:
+        slack_bot_config = get_slack_bot_config_for_channel(
+            channel_name=channel_name, db_session=db_session
+        )
 
-            # Be careful about this default, don't want to accidentally spam every channel
-            # Users should be able to DM slack bot in their private channels though
-            if (
-                slack_bot_config is None
-                and not respond_every_channel
-                # Can't have configs for DMs so don't toss them out
-                and not is_dm
-                # If /DanswerBot (is_bot_msg) or @DanswerBot (bypass_filters)
-                # always respond with the default configs
-                and not (details.is_bot_msg or details.bypass_filters)
-            ):
-                return
+        # Be careful about this default, don't want to accidentally spam every channel
+        # Users should be able to DM slack bot in their private channels though
+        if (
+            slack_bot_config is None
+            and not respond_every_channel
+            # Can't have configs for DMs so don't toss them out
+            and not is_dm
+            # If /DanswerBot (is_bot_msg) or @DanswerBot (bypass_filters)
+            # always respond with the default configs
+            and not (details.is_bot_msg or details.bypass_filters)
+        ):
+            return
 
-            follow_up = bool(
-                slack_bot_config
-                and slack_bot_config.channel_config
-                and slack_bot_config.channel_config.get("follow_up_tags") is not None
-            )
-            feedback_reminder_id = schedule_feedback_reminder(
-                details=details, client=client.web_client, include_followup=follow_up
-            )
+        follow_up = bool(
+            slack_bot_config
+            and slack_bot_config.channel_config
+            and slack_bot_config.channel_config.get("follow_up_tags") is not None
+        )
+        feedback_reminder_id = schedule_feedback_reminder(
+            details=details, client=client.web_client, include_followup=follow_up
+        )
 
-            failed = handle_message(
-                message_info=details,
-                slack_bot_config=slack_bot_config,
-                client=client.web_client,
-                feedback_reminder_id=feedback_reminder_id,
-                tenant_id=client.tenant_id,
-            )
+        failed = handle_message(
+            message_info=details,
+            slack_bot_config=slack_bot_config,
+            client=client.web_client,
+            feedback_reminder_id=feedback_reminder_id,
+        )
 
-            if failed:
-                if feedback_reminder_id:
-                    remove_scheduled_feedback_reminder(
-                        client=client.web_client,
-                        channel=details.sender,
-                        msg_id=feedback_reminder_id,
-                    )
-                # Skipping answering due to pre-filtering is not considered a failure
-                if notify_no_answer:
-                    apologize_for_fail(details, client)
-    finally:
-        if client.tenant_id:
-            CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+        if failed:
+            if feedback_reminder_id:
+                remove_scheduled_feedback_reminder(
+                    client=client.web_client,
+                    channel=details.sender,
+                    msg_id=feedback_reminder_id,
+                )
+            # Skipping answering due to pre-filtering is not considered a failure
+            if notify_no_answer:
+                apologize_for_fail(details, client)
 
 
-def acknowledge_message(req: SocketModeRequest, client: TenantSocketModeClient) -> None:
+def acknowledge_message(req: SocketModeRequest, client: SocketModeClient) -> None:
     response = SocketModeResponse(envelope_id=req.envelope_id)
     client.send_socket_mode_response(response)
 
 
-def action_routing(req: SocketModeRequest, client: TenantSocketModeClient) -> None:
+def action_routing(req: SocketModeRequest, client: SocketModeClient) -> None:
     if actions := req.payload.get("actions"):
         action = cast(dict[str, Any], actions[0])
 
@@ -440,13 +429,13 @@ def action_routing(req: SocketModeRequest, client: TenantSocketModeClient) -> No
             return handle_generate_answer_button(req, client)
 
 
-def view_routing(req: SocketModeRequest, client: TenantSocketModeClient) -> None:
+def view_routing(req: SocketModeRequest, client: SocketModeClient) -> None:
     if view := req.payload.get("view"):
         if view["callback_id"] == VIEW_DOC_FEEDBACK_ID:
             return process_feedback(req, client)
 
 
-def process_slack_event(client: TenantSocketModeClient, req: SocketModeRequest) -> None:
+def process_slack_event(client: SocketModeClient, req: SocketModeRequest) -> None:
     # Always respond right away, if Slack doesn't receive these frequently enough
     # it will assume the Bot is DEAD!!! :(
     acknowledge_message(req, client)
@@ -464,24 +453,21 @@ def process_slack_event(client: TenantSocketModeClient, req: SocketModeRequest)
         logger.error(f"Slack request payload: {req.payload}")
 
 
-def _get_socket_client(
-    slack_bot_tokens: SlackBotTokens, tenant_id: str | None
-) -> TenantSocketModeClient:
+def _get_socket_client(slack_bot_tokens: SlackBotTokens) -> SocketModeClient:
     # For more info on how to set this up, checkout the docs:
     # https://docs.danswer.dev/slack_bot_setup
-    return TenantSocketModeClient(
+    return SocketModeClient(
         # This app-level token will be used only for establishing a connection
         app_token=slack_bot_tokens.app_token,
         web_client=WebClient(token=slack_bot_tokens.bot_token),
-        tenant_id=tenant_id,
     )
 
 
-def _initialize_socket_client(socket_client: TenantSocketModeClient) -> None:
+def _initialize_socket_client(socket_client: SocketModeClient) -> None:
     socket_client.socket_mode_request_listeners.append(process_slack_event)  # type: ignore
 
     # Establish a WebSocket connection to the Socket Mode servers
-    logger.notice(f"Listening for messages from Slack {socket_client.tenant_id }...")
+    logger.notice("Listening for messages from Slack...")
     socket_client.connect()
 
 
@@ -495,8 +481,8 @@ def _initialize_socket_client(socket_client: TenantSocketModeClient) -> None:
 # NOTE: we are using Web Sockets so that you can run this from within a firewalled VPC
 # without issue.
 if __name__ == "__main__":
-    slack_bot_tokens: dict[str | None, SlackBotTokens] = {}
-    socket_clients: dict[str | None, TenantSocketModeClient] = {}
+    slack_bot_tokens: SlackBotTokens | None = None
+    socket_client: SocketModeClient | None = None
 
     set_is_ee_based_on_env_variable()
 
@@ -505,66 +491,46 @@ if __name__ == "__main__":
 
     while True:
         try:
-            tenant_ids = get_all_tenant_ids()  # Function to retrieve all tenant IDs
+            latest_slack_bot_tokens = fetch_tokens()
 
-            for tenant_id in tenant_ids:
-                with get_session_with_tenant(tenant_id) as db_session:
-                    try:
-                        token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id or "public")
-                        latest_slack_bot_tokens = fetch_tokens()
-                        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+            if latest_slack_bot_tokens != slack_bot_tokens:
+                if slack_bot_tokens is not None:
+                    logger.notice("Slack Bot tokens have changed - reconnecting")
+                else:
+                    # This happens on the very first time the listener process comes up
+                    # or the tokens have updated (set up for the first time)
+                    with Session(get_sqlalchemy_engine()) as db_session:
+                        search_settings = get_current_search_settings(db_session)
+                        embedding_model = EmbeddingModel.from_db_model(
+                            search_settings=search_settings,
+                            server_host=MODEL_SERVER_HOST,
+                            server_port=MODEL_SERVER_PORT,
+                        )
 
-                        if (
-                            tenant_id not in slack_bot_tokens
-                            or latest_slack_bot_tokens != slack_bot_tokens[tenant_id]
-                        ):
-                            if tenant_id in slack_bot_tokens:
-                                logger.notice(
-                                    f"Slack Bot tokens have changed for tenant {tenant_id} - reconnecting"
-                                )
-                            else:
-                                # Initial setup for this tenant
-                                search_settings = get_current_search_settings(
-                                    db_session
-                                )
-                                embedding_model = EmbeddingModel.from_db_model(
-                                    search_settings=search_settings,
-                                    server_host=MODEL_SERVER_HOST,
-                                    server_port=MODEL_SERVER_PORT,
-                                )
-                                warm_up_bi_encoder(embedding_model=embedding_model)
+                        warm_up_bi_encoder(
+                            embedding_model=embedding_model,
+                        )
 
-                            slack_bot_tokens[tenant_id] = latest_slack_bot_tokens
+                slack_bot_tokens = latest_slack_bot_tokens
+                # potentially may cause a message to be dropped, but it is complicated
+                # to avoid + (1) if the user is changing tokens, they are likely okay with some
+                # "migration downtime" and (2) if a single message is lost it is okay
+                # as this should be a very rare occurrence
+                if socket_client:
+                    socket_client.close()
 
-                            # potentially may cause a message to be dropped, but it is complicated
-                            # to avoid + (1) if the user is changing tokens, they are likely okay with some
-                            # "migration downtime" and (2) if a single message is lost it is okay
-                            # as this should be a very rare occurrence
-                            if tenant_id in socket_clients:
-                                socket_clients[tenant_id].close()
+                socket_client = _get_socket_client(slack_bot_tokens)
+                _initialize_socket_client(socket_client)
 
-                            socket_client = _get_socket_client(
-                                latest_slack_bot_tokens, tenant_id
-                            )
-
-                            # Initialize socket client for this tenant. Each tenant has its own
-                            # socket client, allowing for multiple concurrent connections (one
-                            # per tenant) with the tenant ID wrapped in the socket model client.
-                            # Each `connect` stores websocket connection in a separate thread.
-                            _initialize_socket_client(socket_client)
-
-                            socket_clients[tenant_id] = socket_client
-
-                    except KvKeyNotFoundError:
-                        logger.debug(f"Missing Slack Bot tokens for tenant {tenant_id}")
-                        if tenant_id in socket_clients:
-                            socket_clients[tenant_id].disconnect()
-                            del socket_clients[tenant_id]
-                            del slack_bot_tokens[tenant_id]
-
-            # Wait before checking for updates
+            # Let the handlers run in the background + re-check for token updates every 60 seconds
             Event().wait(timeout=60)
-
-        except Exception:
-            logger.exception("An error occurred outside of main event loop")
+        except KvKeyNotFoundError:
+            # try again every 30 seconds. This is needed since the user may add tokens
+            # via the UI at any point in the programs lifecycle - if we just allow it to
+            # fail, then the user will need to restart the containers after adding tokens
+            logger.debug(
+                "Missing Slack Bot tokens - waiting 60 seconds and trying again"
+            )
+            if socket_client:
+                socket_client.disconnect()
             time.sleep(60)

backend/danswer/danswerbot/slack/utils.py

@@ -12,7 +12,7 @@ from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
 from slack_sdk.models.blocks import Block
 from slack_sdk.models.metadata import Metadata
-from slack_sdk.socket_mode import SocketModeClient
+from sqlalchemy.orm import Session
 
 from danswer.configs.app_configs import DISABLE_TELEMETRY
 from danswer.configs.constants import ID_SEPARATOR
@@ -31,7 +31,7 @@ from danswer.connectors.slack.utils import make_slack_api_rate_limited
 from danswer.connectors.slack.utils import SlackTextCleaner
 from danswer.danswerbot.slack.constants import FeedbackVisibility
 from danswer.danswerbot.slack.tokens import fetch_tokens
-from danswer.db.engine import get_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_engine
 from danswer.db.users import get_user_by_email
 from danswer.llm.exceptions import GenAIDisabledException
 from danswer.llm.factory import get_default_llms
@@ -489,9 +489,7 @@ def read_slack_thread(
     return thread_messages
 
 
-def slack_usage_report(
-    action: str, sender_id: str | None, client: WebClient, tenant_id: str | None
-) -> None:
+def slack_usage_report(action: str, sender_id: str | None, client: WebClient) -> None:
     if DISABLE_TELEMETRY:
         return
 
@@ -503,7 +501,7 @@ def slack_usage_report(
         logger.warning("Unable to find sender email")
 
     if sender_email is not None:
-        with get_session_with_tenant(tenant_id) as db_session:
+        with Session(get_sqlalchemy_engine()) as db_session:
             danswer_user = get_user_by_email(email=sender_email, db_session=db_session)
 
     optional_telemetry(
@@ -579,9 +577,3 @@ def get_feedback_visibility() -> FeedbackVisibility:
         return FeedbackVisibility(DANSWER_BOT_FEEDBACK_VISIBILITY.lower())
     except ValueError:
         return FeedbackVisibility.PRIVATE
-
-
-class TenantSocketModeClient(SocketModeClient):
-    def __init__(self, tenant_id: str | None, *args: Any, **kwargs: Any):
-        super().__init__(*args, **kwargs)
-        self.tenant_id = tenant_id

backend/danswer/db/auth.py

@@ -10,12 +10,10 @@ from fastapi_users_db_sqlalchemy.access_token import SQLAlchemyAccessTokenDataba
 from sqlalchemy import func
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.future import select
-from sqlalchemy.orm import Session
 
-from danswer.auth.invited_users import get_invited_users
 from danswer.auth.schemas import UserRole
 from danswer.db.engine import get_async_session
-from danswer.db.engine import get_async_session_with_tenant
+from danswer.db.engine import get_sqlalchemy_async_engine
 from danswer.db.models import AccessToken
 from danswer.db.models import OAuthAccount
 from danswer.db.models import User
@@ -35,20 +33,10 @@ def get_default_admin_user_emails() -> list[str]:
     return get_default_admin_user_emails_fn()
 
 
-def get_total_users(db_session: Session) -> int:
-    """
-    Returns the total number of users in the system.
-    This is the sum of users and invited users.
-    """
-    user_count = db_session.query(User).count()
-    invited_users = len(get_invited_users())
-    return user_count + invited_users
-
-
 async def get_user_count() -> int:
-    async with get_async_session_with_tenant() as session:
+    async with AsyncSession(get_sqlalchemy_async_engine()) as asession:
         stmt = select(func.count(User.id))
-        result = await session.execute(stmt)
+        result = await asession.execute(stmt)
         user_count = result.scalar()
         if user_count is None:
             raise RuntimeError("Was not able to fetch the user count.")