minor clarity

* moderate component fixes

* nit

* nit

* update colors

* k

.github/pull_request_template.md

@@ -6,20 +6,24 @@
 [Describe the tests you ran to verify your changes]
 
 
-## Accepted Risk
-[Any know risks or failure modes to point out to reviewers]
+## Accepted Risk (provide if relevant)
+N/A
 
 
-## Related Issue(s)
-[If applicable, link to the issue(s) this PR addresses]
+## Related Issue(s) (provide if relevant)
+N/A
 
 
-## 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
+## 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)

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

@@ -3,61 +3,61 @@ name: Build and Push Backend Image on Tag
 on:
   push:
     tags:
-      - '*'
+      - "*"
 
 env:
-  REGISTRY_IMAGE: danswer/danswer-backend
+  REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'danswer/danswer-backend-cloud' || 'danswer/danswer-backend' }}
   LATEST_TAG: ${{ contains(github.ref_name, 'latest') }}
-  
+
 jobs:
   build-and-push:
-    # TODO: investigate a matrix build like the web container 
+    # TODO: investigate a matrix build like the web container
     # See https://runs-on.com/runners/linux/
-    runs-on: [runs-on,runner=8cpu-linux-x64,"run-id=${{ github.run_id }}"]
+    runs-on: [runs-on, runner=8cpu-linux-x64, "run-id=${{ github.run_id }}"]
 
     steps:
-    - name: Checkout code
-      uses: actions/checkout@v4
+      - name: Checkout code
+        uses: actions/checkout@v4
 
-    - name: Set up Docker Buildx
-      uses: docker/setup-buildx-action@v3
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
 
-    - name: Login to Docker Hub
-      uses: docker/login-action@v3
-      with:
-        username: ${{ secrets.DOCKER_USERNAME }}
-        password: ${{ secrets.DOCKER_TOKEN }}
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_TOKEN }}
 
-    - name: Install build-essential
-      run: |
-        sudo apt-get update
-        sudo apt-get install -y build-essential
-          
-    - name: Backend Image Docker Build and Push
-      uses: docker/build-push-action@v5
-      with:
-        context: ./backend
-        file: ./backend/Dockerfile
-        platforms: linux/amd64,linux/arm64
-        push: true
-        tags: |
-          ${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
-          ${{ env.LATEST_TAG == 'true' && format('{0}:latest', env.REGISTRY_IMAGE) || '' }}
-        build-args: |
-          DANSWER_VERSION=${{ github.ref_name }}
+      - name: Install build-essential
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y build-essential
 
-    # trivy has their own rate limiting issues causing this action to flake
-    # we worked around it by hardcoding to different db repos in env
-    # can re-enable when they figure it out
-    # https://github.com/aquasecurity/trivy/discussions/7538
-    # https://github.com/aquasecurity/trivy-action/issues/389
-    - name: Run Trivy vulnerability scanner
-      uses: aquasecurity/trivy-action@master
-      env:
-        TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
-        TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
-      with:
-        # To run locally: trivy image --severity HIGH,CRITICAL danswer/danswer-backend
-        image-ref: docker.io/${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
-        severity: 'CRITICAL,HIGH'
-        trivyignores: ./backend/.trivyignore
+      - name: Backend Image Docker Build and Push
+        uses: docker/build-push-action@v5
+        with:
+          context: ./backend
+          file: ./backend/Dockerfile
+          platforms: linux/amd64,linux/arm64
+          push: true
+          tags: |
+            ${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
+            ${{ env.LATEST_TAG == 'true' && format('{0}:latest', env.REGISTRY_IMAGE) || '' }}
+          build-args: |
+            DANSWER_VERSION=${{ github.ref_name }}
+
+      # trivy has their own rate limiting issues causing this action to flake
+      # we worked around it by hardcoding to different db repos in env
+      # can re-enable when they figure it out
+      # https://github.com/aquasecurity/trivy/discussions/7538
+      # https://github.com/aquasecurity/trivy-action/issues/389
+      - name: Run Trivy vulnerability scanner
+        uses: aquasecurity/trivy-action@master
+        env:
+          TRIVY_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-db:2"
+          TRIVY_JAVA_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-java-db:1"
+        with:
+          # To run locally: trivy image --severity HIGH,CRITICAL danswer/danswer-backend
+          image-ref: docker.io/${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
+          severity: "CRITICAL,HIGH"
+          trivyignores: ./backend/.trivyignore

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

@@ -0,0 +1,136 @@
+name: Build and Push Cloud Web Image on Tag
+# Identical to the web container build, but with correct image tag and build args
+
+on:
+  push:
+    tags:
+      - "*"
+
+env:
+  REGISTRY_IMAGE: danswer/danswer-web-server-cloud
+  LATEST_TAG: ${{ contains(github.ref_name, 'latest') }}
+
+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 }}
+    strategy:
+      fail-fast: false
+      matrix:
+        platform:
+          - linux/amd64
+          - linux/arm64
+
+    steps:
+      - name: Prepare
+        run: |
+          platform=${{ matrix.platform }}
+          echo "PLATFORM_PAIR=${platform//\//-}" >> $GITHUB_ENV
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Docker meta
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ${{ env.REGISTRY_IMAGE }}
+          tags: |
+            type=raw,value=${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
+            type=raw,value=${{ env.LATEST_TAG == 'true' && format('{0}:latest', env.REGISTRY_IMAGE) || '' }}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_TOKEN }}
+
+      - name: Build and push by digest
+        id: build
+        uses: docker/build-push-action@v5
+        with:
+          context: ./web
+          file: ./web/Dockerfile
+          platforms: ${{ matrix.platform }}
+          push: true
+          build-args: |
+            DANSWER_VERSION=${{ github.ref_name }}
+            NEXT_PUBLIC_CLOUD_ENABLED=true
+            NEXT_PUBLIC_POSTHOG_KEY=${{ secrets.POSTHOG_KEY }}
+            NEXT_PUBLIC_POSTHOG_HOST=${{ secrets.POSTHOG_HOST }}
+            NEXT_PUBLIC_SENTRY_DSN=${{ secrets.SENTRY_DSN }}
+          # needed due to weird interactions with the builds for different platforms
+          no-cache: true
+          labels: ${{ steps.meta.outputs.labels }}
+          outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true
+
+      - name: Export digest
+        run: |
+          mkdir -p /tmp/digests
+          digest="${{ steps.build.outputs.digest }}"
+          touch "/tmp/digests/${digest#sha256:}"
+
+      - name: Upload digest
+        uses: actions/upload-artifact@v4
+        with:
+          name: digests-${{ env.PLATFORM_PAIR }}
+          path: /tmp/digests/*
+          if-no-files-found: error
+          retention-days: 1
+
+  merge:
+    runs-on: ubuntu-latest
+    needs:
+      - build
+    steps:
+      - name: Download digests
+        uses: actions/download-artifact@v4
+        with:
+          path: /tmp/digests
+          pattern: digests-*
+          merge-multiple: true
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Docker meta
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ${{ env.REGISTRY_IMAGE }}
+
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_TOKEN }}
+
+      - name: Create manifest list and push
+        working-directory: /tmp/digests
+        run: |
+          docker buildx imagetools create $(jq -cr '.tags | map("-t " + .) | join(" ")' <<< "$DOCKER_METADATA_OUTPUT_JSON") \
+            $(printf '${{ env.REGISTRY_IMAGE }}@sha256:%s ' *)
+
+      - name: Inspect image
+        run: |
+          docker buildx imagetools inspect ${{ env.REGISTRY_IMAGE }}:${{ steps.meta.outputs.version }}
+
+      # trivy has their own rate limiting issues causing this action to flake
+      # we worked around it by hardcoding to different db repos in env
+      # can re-enable when they figure it out
+      # https://github.com/aquasecurity/trivy/discussions/7538
+      # https://github.com/aquasecurity/trivy-action/issues/389
+      - name: Run Trivy vulnerability scanner
+        uses: aquasecurity/trivy-action@master
+        env:
+          TRIVY_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-db:2"
+          TRIVY_JAVA_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-java-db:1"
+        with:
+          image-ref: docker.io/${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
+          severity: "CRITICAL,HIGH"

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

@@ -3,53 +3,53 @@ name: Build and Push Model Server Image on Tag
 on:
   push:
     tags:
-      - '*'
+      - "*"
 
 env:
-  REGISTRY_IMAGE: danswer/danswer-model-server
+  REGISTRY_IMAGE: ${{ contains(github.ref_name, 'cloud') && 'danswer/danswer-model-server-cloud' || 'danswer/danswer-model-server' }}
   LATEST_TAG: ${{ contains(github.ref_name, 'latest') }}
-  
+
 jobs:
   build-and-push:
     # See https://runs-on.com/runners/linux/
-    runs-on: [runs-on,runner=8cpu-linux-x64,"run-id=${{ github.run_id }}"]
+    runs-on: [runs-on, runner=8cpu-linux-x64, "run-id=${{ github.run_id }}"]
 
     steps:
-    - name: Checkout code
-      uses: actions/checkout@v4
+      - name: Checkout code
+        uses: actions/checkout@v4
 
-    - name: Set up Docker Buildx
-      uses: docker/setup-buildx-action@v3
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
 
-    - name: Login to Docker Hub
-      uses: docker/login-action@v3
-      with:
-        username: ${{ secrets.DOCKER_USERNAME }}
-        password: ${{ secrets.DOCKER_TOKEN }}
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKER_USERNAME }}
+          password: ${{ secrets.DOCKER_TOKEN }}
 
-    - name: Model Server Image Docker Build and Push
-      uses: docker/build-push-action@v5
-      with:
-        context: ./backend
-        file: ./backend/Dockerfile.model_server
-        platforms: linux/amd64,linux/arm64
-        push: true
-        tags: |
-          ${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
-          ${{ env.LATEST_TAG == 'true' && format('{0}:latest', env.REGISTRY_IMAGE) || '' }}
-        build-args: |
-          DANSWER_VERSION=${{ github.ref_name }}
+      - name: Model Server Image Docker Build and Push
+        uses: docker/build-push-action@v5
+        with:
+          context: ./backend
+          file: ./backend/Dockerfile.model_server
+          platforms: linux/amd64,linux/arm64
+          push: true
+          tags: |
+            ${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
+            ${{ env.LATEST_TAG == 'true' && format('{0}:latest', env.REGISTRY_IMAGE) || '' }}
+          build-args: |
+            DANSWER_VERSION=${{ github.ref_name }}
 
-    # trivy has their own rate limiting issues causing this action to flake
-    # we worked around it by hardcoding to different db repos in env
-    # can re-enable when they figure it out
-    # https://github.com/aquasecurity/trivy/discussions/7538
-    # https://github.com/aquasecurity/trivy-action/issues/389
-    - name: Run Trivy vulnerability scanner
-      uses: aquasecurity/trivy-action@master
-      env:
-        TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
-        TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
-      with:
-        image-ref: docker.io/danswer/danswer-model-server:${{ github.ref_name }}
-        severity: 'CRITICAL,HIGH'
+      # trivy has their own rate limiting issues causing this action to flake
+      # we worked around it by hardcoding to different db repos in env
+      # can re-enable when they figure it out
+      # https://github.com/aquasecurity/trivy/discussions/7538
+      # https://github.com/aquasecurity/trivy-action/issues/389
+      - name: Run Trivy vulnerability scanner
+        uses: aquasecurity/trivy-action@master
+        env:
+          TRIVY_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-db:2"
+          TRIVY_JAVA_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-java-db:1"
+        with:
+          image-ref: docker.io/danswer/danswer-model-server:${{ github.ref_name }}
+          severity: "CRITICAL,HIGH"

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

@@ -0,0 +1,23 @@
+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/nightly-scan-licenses.yml

@@ -0,0 +1,76 @@
+# Scan for problematic software licenses
+
+# trivy has their own rate limiting issues causing this action to flake
+# we worked around it by hardcoding to different db repos in env
+# can re-enable when they figure it out
+# https://github.com/aquasecurity/trivy/discussions/7538
+# https://github.com/aquasecurity/trivy-action/issues/389
+
+name: 'Nightly - Scan licenses'
+on:
+#   schedule:
+#     - cron: '0 14 * * *'  # Runs every day at 6 AM PST / 7 AM PDT / 2 PM UTC
+  workflow_dispatch:  # Allows manual triggering
+
+permissions:
+  actions: read
+  contents: read
+  security-events: write
+  
+jobs:
+  scan-licenses:
+    # See https://runs-on.com/runners/linux/
+    runs-on: [runs-on,runner=2cpu-linux-x64,"run-id=${{ github.run_id }}"]
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+          cache-dependency-path: |
+            backend/requirements/default.txt
+            backend/requirements/dev.txt
+            backend/requirements/model_server.txt
+      
+      - name: Get explicit and transitive dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install --retries 5 --timeout 30 -r backend/requirements/default.txt
+          pip install --retries 5 --timeout 30 -r backend/requirements/dev.txt
+          pip install --retries 5 --timeout 30 -r backend/requirements/model_server.txt
+          pip freeze > requirements-all.txt
+                    
+      - name: Check python
+        id: license_check_report
+        uses: pilosus/action-pip-license-checker@v2
+        with:
+          requirements: 'requirements-all.txt'
+          fail: 'Copyleft'
+          exclude: '(?i)^(pylint|aio[-_]*).*'
+          
+      - name: Print report
+        if: ${{ always() }}
+        run: echo "${{ steps.license_check_report.outputs.report }}"
+      
+      - name: Install npm dependencies
+        working-directory: ./web
+        run: npm ci
+        
+      - name: Run Trivy vulnerability scanner in repo mode
+        uses: aquasecurity/trivy-action@0.28.0
+        with:
+          scan-type: fs
+          scanners: license
+          format: table
+#           format: sarif
+#           output: trivy-results.sarif
+          severity: HIGH,CRITICAL
+
+#       - name: Upload Trivy scan results to GitHub Security tab
+#         uses: github/codeql-action/upload-sarif@v3
+#         with:
+#           sarif_file: trivy-results.sarif

.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,7 +85,58 @@ 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
 
-      - name: Start Docker containers
+      # 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 
         run: |
           cd deployment/docker_compose
           ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true \
@@ -130,7 +181,7 @@ jobs:
           done
           echo "Finished waiting for service."
 
-      - name: Run integration tests
+      - name: Run Standard Integration Tests
         run: |
           echo "Running integration tests..."
           docker run --rm --network danswer-stack_default \
@@ -145,7 +196,8 @@ 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
+            danswer/danswer-integration:test \
+            /app/tests/integration/tests
         continue-on-error: true
         id: run_tests
 
@@ -158,12 +210,18 @@ jobs:
             echo "All integration tests passed successfully."
           fi
 
+      # save before stopping the containers so the logs can be captured
       - name: Save Docker logs
         if: success() || failure()
         run: |
           cd deployment/docker_compose
           docker compose -f docker-compose.dev.yml -p danswer-stack logs > docker-compose.log
           mv docker-compose.log ${{ github.workspace }}/docker-compose.log
+
+      - name: Stop Docker containers
+        run: |
+          cd deployment/docker_compose
+          docker compose -f docker-compose.dev.yml -p danswer-stack down -v
       
       - name: Upload logs
         if: success() || failure()

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

@@ -0,0 +1,124 @@
+name: Backport on Merge
+
+# Note this workflow does not trigger the builds, be sure to manually tag the branches to trigger the builds
+
+on:
+  pull_request:
+    types: [closed]  # Later we check for merge so only PRs that go in can get backported
+
+permissions:
+  contents: write
+  actions: write
+
+jobs:
+  backport:
+    if: github.event.pull_request.merged == true
+    runs-on: ubuntu-latest
+    env:
+      GITHUB_TOKEN: ${{ secrets.YUHONG_GH_ACTIONS }}
+    steps:
+      - name: Checkout code
+        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"
+          git fetch --prune
+      
+      - 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)
+          
+          # Handle case where no beta tags exist
+          if [[ -z "$LATEST_BETA_TAG" ]]; then
+            NEW_BETA_TAG="v1.0.0-beta.1"
+          else
+            NEW_BETA_TAG=$(echo $LATEST_BETA_TAG | awk -F '[.-]' '{print $1 "." $2 "." $3 "-beta." ($NF+1)}')
+          fi
+          
+          # 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 }}"
+          
+          # Echo the merge commit SHA
+          echo "Merge commit SHA: ${{ github.event.pull_request.merge_commit_sha }}"
+          
+          # Fetch all history for all branches and tags
+          git fetch --prune
+
+          # Reset and prepare the beta branch
+          git checkout ${{ steps.list-branches.outputs.beta }}
+          echo "Last 5 commits on beta branch:"
+          git log -n 5 --pretty=format:"%H"
+          echo ""  # Newline for formatting
+
+          # 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 branch/tag
+          git tag ${{ steps.list-branches.outputs.new_beta_tag }}
+          # Push the changes and tag to the beta branch using PAT
+          git push origin ${{ steps.list-branches.outputs.beta }}
+          git push origin ${{ steps.list-branches.outputs.new_beta_tag }}
+
+          # Reset and prepare the stable branch
+          git checkout ${{ steps.list-branches.outputs.stable }}
+          echo "Last 5 commits on stable branch:"
+          git log -n 5 --pretty=format:"%H"
+          echo ""  # Newline for formatting
+          
+          # 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 branch/tag
+          git tag ${{ steps.list-branches.outputs.new_stable_tag }}
+          # Push the changes and tag to the stable branch using PAT
+          git push origin ${{ steps.list-branches.outputs.stable }}
+          git push origin ${{ steps.list-branches.outputs.new_stable_tag }}

.github/workflows/pr-helm-chart-testing.yml

@@ -1,24 +1,20 @@
-# This workflow is intentionally disabled while we're still working on it
-# It's close to ready, but a race condition needs to be fixed with
-# API server and Vespa startup, and it needs to have a way to build/test against
-# local containers
-
 name: Helm - Lint and Test Charts
 
 on:
   merge_group:
   pull_request:
     branches: [ main ]
-
+  workflow_dispatch:  # Allows manual triggering
+  
 jobs:
-  lint-test:
+  helm-chart-check:
     # See https://runs-on.com/runners/linux/
     runs-on: [runs-on,runner=8cpu-linux-x64,hdd=256,"run-id=${{ github.run_id }}"]
 
     # fetch-depth 0 is required for helm/chart-testing-action
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v4
       with:
         fetch-depth: 0
         
@@ -28,7 +24,7 @@ jobs:
         version: v3.14.4
       
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.11'
         cache: 'pip'
@@ -45,24 +41,31 @@ jobs:
     - name: Set up chart-testing
       uses: helm/chart-testing-action@v2.6.1
 
+    # even though we specify chart-dirs in ct.yaml, it isn't used by ct for the list-changed command...
     - name: Run chart-testing (list-changed)
       id: list-changed
       run: |
-        changed=$(ct list-changed --target-branch ${{ github.event.repository.default_branch }})
+        echo "default_branch: ${{ github.event.repository.default_branch }}"
+        changed=$(ct list-changed --remote origin --target-branch ${{ github.event.repository.default_branch }} --chart-dirs deployment/helm/charts)
+        echo "list-changed output: $changed"
         if [[ -n "$changed" ]]; then
           echo "changed=true" >> "$GITHUB_OUTPUT"
         fi
 
+    # lint all charts if any changes were detected
     - name: Run chart-testing (lint)
-#       if: steps.list-changed.outputs.changed == 'true'
-      run: ct lint --all --config ct.yaml --target-branch ${{ github.event.repository.default_branch }}
+      if: steps.list-changed.outputs.changed == 'true'
+      run: ct lint --config ct.yaml --all
+      # the following would lint only changed charts, but linting isn't expensive
+      # run: ct lint --config ct.yaml --target-branch ${{ github.event.repository.default_branch }}
 
     - name: Create kind cluster
-#       if: steps.list-changed.outputs.changed == 'true'
+      if: steps.list-changed.outputs.changed == 'true'
       uses: helm/kind-action@v1.10.0
 
     - name: Run chart-testing (install)
-#       if: steps.list-changed.outputs.changed == 'true'
-      run: ct install --all --config ct.yaml
-#       run: ct install --target-branch ${{ github.event.repository.default_branch }}
-      
+      if: steps.list-changed.outputs.changed == 'true'
+      run: ct install --all --helm-extra-set-args="--set=nginx.enabled=false" --debug --config ct.yaml
+      # the following would install only changed charts, but we only have one chart so 
+      # don't worry about that for now
+      # run: ct install --target-branch ${{ github.event.repository.default_branch }}

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

@@ -18,6 +18,11 @@ env:
   # Jira
   JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}
   JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}
+  # Google
+  GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR: ${{ secrets.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR }}
+  GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR: ${{ secrets.GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR }}
+  GOOGLE_GMAIL_SERVICE_ACCOUNT_JSON_STR: ${{ secrets.GOOGLE_GMAIL_SERVICE_ACCOUNT_JSON_STR }}
+  GOOGLE_GMAIL_OAUTH_CREDENTIALS_JSON_STR: ${{ secrets.GOOGLE_GMAIL_OAUTH_CREDENTIALS_JSON_STR }}
 
 jobs:
   connectors-check:

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

@@ -15,7 +15,7 @@ env:
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
 
 jobs:
-  connectors-check:
+  model-check:
     # See https://runs-on.com/runners/linux/
     runs-on: [runs-on,runner=8cpu-linux-x64,"run-id=${{ github.run_id }}"]
 

.vscode/launch.template.jsonc

@@ -6,19 +6,69 @@
     // 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",
-                "Indexing",
-                "Background Jobs",
-                "Slack Bot"
-            ]
-        }
+                "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",
+			 }
+		}
     ],
     "configurations": [
+		{
+	    	// Dummy entry used to label the group
+			"name": "--- Individual ---",
+			"type": "node",
+			"request": "launch",
+			"presentation": {
+				"group": "2",
+				"order": 0
+			}
+		},
         {
             "name": "Web Server",
             "type": "node",
@@ -29,7 +79,11 @@
             "runtimeArgs": [
                 "run", "dev"
             ],
-            "console": "integratedTerminal"
+            "presentation": {
+				 "group": "2",
+			 },
+            "console": "integratedTerminal",
+            "consoleTitle": "Web Server Console"
         },
         {
             "name": "Model Server",
@@ -48,7 +102,11 @@
                 "--reload",
                 "--port",
                 "9000"
-            ]
+            ],
+            "presentation": {
+				 "group": "2",
+			},
+            "consoleTitle": "Model Server Console"
         },
         {
             "name": "API Server",
@@ -68,43 +126,13 @@
                 "--reload",
                 "--port",
                 "8080"
-            ]
+            ],
+            "presentation": {
+				 "group": "2",
+			 },
+            "consoleTitle": "API Server Console"
         },
-        {
-            "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,
+        // For the listener 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",
@@ -118,7 +146,151 @@
                 "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",
@@ -137,8 +309,22 @@
                 "-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",
@@ -147,7 +333,27 @@
             "runtimeArgs": ["${workspaceFolder}/backend/scripts/restart_containers.sh"],
             "cwd": "${workspaceFolder}",
             "console": "integratedTerminal",
-            "stopOnEntry": true
-        }
+            "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": "."
+            },
+        },
     ]
 }

README.md

@@ -1,4 +1,5 @@
 <!-- DANSWER_METADATA={"link": "https://github.com/danswer-ai/danswer/blob/main/README.md"} -->
+<a name="readme-top"></a>
 
 <h2 align="center">
 <a href="https://www.danswer.ai/"> <img width="50%" src="https://github.com/danswer-owners/danswer/blob/1fabd9372d66cd54238847197c33f091a724803b/DanswerWithName.png?raw=true)" /></a>
@@ -68,7 +69,7 @@ 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.
-* Multi-Model model support, chat with images, video etc.
+* Multimodal 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.
@@ -127,3 +128,19 @@ To try the Danswer Enterprise Edition:
 
 ## 💡 Contributing
 Looking to contribute? Please check out the [Contribution Guide](CONTRIBUTING.md) for more details.
+
+## ⭐Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=danswer-ai/danswer&type=Date)](https://star-history.com/#danswer-ai/danswer&Date)
+
+## ✨Contributors
+
+<a href="https://github.com/aryn-ai/sycamore/graphs/contributors">
+  <img alt="contributors" src="https://contrib.rocks/image?repo=danswer-ai/danswer"/>
+</a>
+
+<p align="right" style="font-size: 14px; color: #555; margin-top: 20px;">
+    <a href="#readme-top" style="text-decoration: none; color: #007bff; font-weight: bold;">
+        ↑ Back to Top ↑
+    </a>
+</p>

backend/Dockerfile

@@ -12,7 +12,6 @@ ARG DANSWER_VERSION=0.8-dev
 ENV DANSWER_VERSION=${DANSWER_VERSION} \
     DANSWER_RUNNING_IN_DOCKER="true"
 
-ARG CA_CERT_CONTENT=""
 
 RUN echo "DANSWER_VERSION: ${DANSWER_VERSION}"
 # Install system dependencies
@@ -39,15 +38,6 @@ RUN apt-get update && \
     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
@@ -87,7 +77,6 @@ RUN apt-get update && \
 RUN python -c "from tokenizers import Tokenizer; \
 Tokenizer.from_pretrained('nomic-ai/nomic-embed-text-v1')"
 
-
 # Pre-downloading NLTK for setups with limited egress
 RUN python -c "import nltk; \
 nltk.download('stopwords', quiet=True); \

backend/alembic/env.py

@@ -9,11 +9,12 @@ 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 shared_configs.configs import MULTI_TENANT
 from danswer.db.engine import build_connection_string
 from danswer.db.models import Base
 from celery.backends.database.session import ResultModelBase  # type: ignore
-from danswer.background.celery.celery_app import get_all_tenant_ids
+from danswer.db.engine import get_all_tenant_ids
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 
 # Alembic Config object
 config = context.config
@@ -57,11 +58,15 @@ def get_schema_options() -> tuple[str, bool, bool]:
             if "=" in pair:
                 key, value = pair.split("=", 1)
                 x_args[key.strip()] = value.strip()
-    schema_name = x_args.get("schema", "public")
+    schema_name = x_args.get("schema", POSTGRES_DEFAULT_SCHEMA)
     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 == "public":
+    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."

backend/alembic/versions/33cb72ea4d80_single_tool_call_per_message.py

@@ -0,0 +1,50 @@
+"""single tool call per message
+
+Revision ID: 33cb72ea4d80
+Revises: 5b29123cd710
+Create Date: 2024-11-01 12:51:01.535003
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "33cb72ea4d80"
+down_revision = "5b29123cd710"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Step 1: Delete extraneous ToolCall entries
+    # Keep only the ToolCall with the smallest 'id' for each 'message_id'
+    op.execute(
+        sa.text(
+            """
+            DELETE FROM tool_call
+            WHERE id NOT IN (
+                SELECT MIN(id)
+                FROM tool_call
+                WHERE message_id IS NOT NULL
+                GROUP BY message_id
+            );
+        """
+        )
+    )
+
+    # Step 2: Add a unique constraint on message_id
+    op.create_unique_constraint(
+        constraint_name="uq_tool_call_message_id",
+        table_name="tool_call",
+        columns=["message_id"],
+    )
+
+
+def downgrade() -> None:
+    # Step 1: Drop the unique constraint on message_id
+    op.drop_constraint(
+        constraint_name="uq_tool_call_message_id",
+        table_name="tool_call",
+        type_="unique",
+    )

backend/alembic/versions/5b29123cd710_nullable_search_settings_for_historic_.py

@@ -0,0 +1,70 @@
+"""nullable search settings for historic index attempts
+
+Revision ID: 5b29123cd710
+Revises: 949b4a92a401
+Create Date: 2024-10-30 19:37:59.630704
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "5b29123cd710"
+down_revision = "949b4a92a401"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Drop the existing foreign key constraint
+    op.drop_constraint(
+        "fk_index_attempt_search_settings", "index_attempt", type_="foreignkey"
+    )
+
+    # Modify the column to be nullable
+    op.alter_column(
+        "index_attempt", "search_settings_id", existing_type=sa.INTEGER(), nullable=True
+    )
+
+    # Add back the foreign key with ON DELETE SET NULL
+    op.create_foreign_key(
+        "fk_index_attempt_search_settings",
+        "index_attempt",
+        "search_settings",
+        ["search_settings_id"],
+        ["id"],
+        ondelete="SET NULL",
+    )
+
+
+def downgrade() -> None:
+    # Warning: This will delete all index attempts that don't have search settings
+    op.execute(
+        """
+        DELETE FROM index_attempt
+        WHERE search_settings_id IS NULL
+    """
+    )
+
+    # Drop foreign key constraint
+    op.drop_constraint(
+        "fk_index_attempt_search_settings", "index_attempt", type_="foreignkey"
+    )
+
+    # Modify the column to be not nullable
+    op.alter_column(
+        "index_attempt",
+        "search_settings_id",
+        existing_type=sa.INTEGER(),
+        nullable=False,
+    )
+
+    # Add back the foreign key without ON DELETE SET NULL
+    op.create_foreign_key(
+        "fk_index_attempt_search_settings",
+        "index_attempt",
+        "search_settings",
+        ["search_settings_id"],
+        ["id"],
+    )

backend/alembic/versions/6756efa39ada_id_uuid_for_chat_session.py

@@ -1,7 +1,9 @@
-"""
+"""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
@@ -12,8 +14,6 @@ branch_labels = None
 depends_on = None
 
 """
-Migrate chat_session and chat_message tables to use UUID primary keys.
-
 This script:
 1. Adds UUID columns to chat_session and chat_message
 2. Populates new columns with UUIDs

backend/alembic/versions/949b4a92a401_remove_rt.py

@@ -0,0 +1,74 @@
+"""remove rt
+
+Revision ID: 949b4a92a401
+Revises: 1b10e1fda030
+Create Date: 2024-10-26 13:06:06.937969
+
+"""
+from alembic import op
+from sqlalchemy.orm import Session
+
+# Import your models and constants
+from danswer.db.models import (
+    Connector,
+    ConnectorCredentialPair,
+    Credential,
+    IndexAttempt,
+)
+from danswer.configs.constants import DocumentSource
+
+
+# revision identifiers, used by Alembic.
+revision = "949b4a92a401"
+down_revision = "1b10e1fda030"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Deletes all RequestTracker connectors and associated data
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    connectors_to_delete = (
+        session.query(Connector)
+        .filter(Connector.source == DocumentSource.REQUESTTRACKER)
+        .all()
+    )
+
+    connector_ids = [connector.id for connector in connectors_to_delete]
+
+    if connector_ids:
+        cc_pairs_to_delete = (
+            session.query(ConnectorCredentialPair)
+            .filter(ConnectorCredentialPair.connector_id.in_(connector_ids))
+            .all()
+        )
+
+        cc_pair_ids = [cc_pair.id for cc_pair in cc_pairs_to_delete]
+
+        if cc_pair_ids:
+            session.query(IndexAttempt).filter(
+                IndexAttempt.connector_credential_pair_id.in_(cc_pair_ids)
+            ).delete(synchronize_session=False)
+
+            session.query(ConnectorCredentialPair).filter(
+                ConnectorCredentialPair.id.in_(cc_pair_ids)
+            ).delete(synchronize_session=False)
+
+        credential_ids = [cc_pair.credential_id for cc_pair in cc_pairs_to_delete]
+        if credential_ids:
+            session.query(Credential).filter(Credential.id.in_(credential_ids)).delete(
+                synchronize_session=False
+            )
+
+        session.query(Connector).filter(Connector.id.in_(connector_ids)).delete(
+            synchronize_session=False
+        )
+
+    session.commit()
+
+
+def downgrade() -> None:
+    # No-op downgrade as we cannot restore deleted data
+    pass

backend/alembic/versions/b082fec533f0_make_last_attempt_status_nullable.py

@@ -31,6 +31,12 @@ 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/b156fa702355_chat_reworked.py

@@ -288,6 +288,15 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
+    # NOTE: you will lose all chat history. This is to satisfy the non-nullable constraints
+    # below
+    op.execute("DELETE FROM chat_feedback")
+    op.execute("DELETE FROM chat_message__search_doc")
+    op.execute("DELETE FROM document_retrieval_feedback")
+    op.execute("DELETE FROM document_retrieval_feedback")
+    op.execute("DELETE FROM chat_message")
+    op.execute("DELETE FROM chat_session")
+
     op.drop_constraint(
         "chat_feedback__chat_message_fk", "chat_feedback", type_="foreignkey"
     )

backend/alembic/versions/b72ed7a5db0e_remove_description_from_starter_messages.py

@@ -0,0 +1,48 @@
+"""remove description from starter messages
+
+Revision ID: b72ed7a5db0e
+Revises: 33cb72ea4d80
+Create Date: 2024-11-03 15:55:28.944408
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "b72ed7a5db0e"
+down_revision = "33cb72ea4d80"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.execute(
+        sa.text(
+            """
+            UPDATE persona
+            SET starter_messages = (
+                SELECT jsonb_agg(elem - 'description')
+                FROM jsonb_array_elements(starter_messages) elem
+            )
+            WHERE starter_messages IS NOT NULL
+              AND jsonb_typeof(starter_messages) = 'array'
+            """
+        )
+    )
+
+
+def downgrade() -> None:
+    op.execute(
+        sa.text(
+            """
+            UPDATE persona
+            SET starter_messages = (
+                SELECT jsonb_agg(elem || '{"description": ""}')
+                FROM jsonb_array_elements(starter_messages) elem
+            )
+            WHERE starter_messages IS NOT NULL
+              AND jsonb_typeof(starter_messages) = 'array'
+            """
+        )
+    )

backend/alembic/versions/c0fd6e4da83a_add_recent_assistants.py

@@ -0,0 +1,29 @@
+"""add recent assistants
+
+Revision ID: c0fd6e4da83a
+Revises: b72ed7a5db0e
+Create Date: 2024-11-03 17:28:54.916618
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "c0fd6e4da83a"
+down_revision = "b72ed7a5db0e"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "user",
+        sa.Column(
+            "recent_assistants", postgresql.JSONB(), server_default="[]", nullable=False
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("user", "recent_assistants")

backend/alembic/versions/c99d76fcd298_add_nullable_to_persona_id_in_chat_.py

@@ -23,6 +23,56 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
+    # Delete chat messages and feedback first since they reference chat sessions
+    # Get chat messages from sessions with null persona_id
+    chat_messages_query = """
+        SELECT id
+        FROM chat_message
+        WHERE chat_session_id IN (
+            SELECT id
+            FROM chat_session
+            WHERE persona_id IS NULL
+        )
+    """
+
+    # Delete dependent records first
+    op.execute(
+        f"""
+        DELETE FROM document_retrieval_feedback
+        WHERE chat_message_id IN (
+            {chat_messages_query}
+        )
+    """
+    )
+    op.execute(
+        f"""
+        DELETE FROM chat_message__search_doc
+        WHERE chat_message_id IN (
+            {chat_messages_query}
+        )
+    """
+    )
+
+    # Delete chat messages
+    op.execute(
+        """
+        DELETE FROM chat_message
+        WHERE chat_session_id IN (
+            SELECT id
+            FROM chat_session
+            WHERE persona_id IS NULL
+        )
+    """
+    )
+
+    # Now we can safely delete the chat sessions
+    op.execute(
+        """
+        DELETE FROM chat_session
+        WHERE persona_id IS NULL
+    """
+    )
+
     op.alter_column(
         "chat_session",
         "persona_id",

backend/danswer/__init__.py

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

backend/danswer/access/models.py

@@ -70,3 +70,12 @@ 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/api_key.py

@@ -0,0 +1,89 @@
+import secrets
+import uuid
+from urllib.parse import quote
+from urllib.parse import unquote
+
+from fastapi import Request
+from passlib.hash import sha256_crypt
+from pydantic import BaseModel
+
+from danswer.auth.schemas import UserRole
+from danswer.configs.app_configs import API_KEY_HASH_ROUNDS
+
+
+_API_KEY_HEADER_NAME = "Authorization"
+# NOTE for others who are curious: In the context of a header, "X-" often refers
+# to non-standard, experimental, or custom headers in HTTP or other protocols. It
+# indicates that the header is not part of the official standards defined by
+# organizations like the Internet Engineering Task Force (IETF).
+_API_KEY_HEADER_ALTERNATIVE_NAME = "X-Danswer-Authorization"
+_BEARER_PREFIX = "Bearer "
+_API_KEY_PREFIX = "dn_"
+_API_KEY_LEN = 192
+
+
+class ApiKeyDescriptor(BaseModel):
+    api_key_id: int
+    api_key_display: str
+    api_key: str | None = None  # only present on initial creation
+    api_key_name: str | None = None
+    api_key_role: UserRole
+
+    user_id: uuid.UUID
+
+
+def generate_api_key(tenant_id: str | None = None) -> str:
+    # For backwards compatibility, if no tenant_id, generate old style key
+    if not tenant_id:
+        return _API_KEY_PREFIX + secrets.token_urlsafe(_API_KEY_LEN)
+
+    encoded_tenant = quote(tenant_id)  # URL encode the tenant ID
+    return f"{_API_KEY_PREFIX}{encoded_tenant}.{secrets.token_urlsafe(_API_KEY_LEN)}"
+
+
+def extract_tenant_from_api_key_header(request: Request) -> str | None:
+    """Extract tenant ID from request. Returns None if auth is disabled or invalid format."""
+    raw_api_key_header = request.headers.get(
+        _API_KEY_HEADER_ALTERNATIVE_NAME
+    ) or request.headers.get(_API_KEY_HEADER_NAME)
+
+    if not raw_api_key_header or not raw_api_key_header.startswith(_BEARER_PREFIX):
+        return None
+
+    api_key = raw_api_key_header[len(_BEARER_PREFIX) :].strip()
+
+    if not api_key.startswith(_API_KEY_PREFIX):
+        return None
+
+    parts = api_key[len(_API_KEY_PREFIX) :].split(".", 1)
+    if len(parts) != 2:
+        return None
+
+    tenant_id = parts[0]
+    return unquote(tenant_id) if tenant_id else None
+
+
+def hash_api_key(api_key: str) -> str:
+    # NOTE: no salt is needed, as the API key is randomly generated
+    # and overlaps are impossible
+    return sha256_crypt.hash(api_key, salt="", rounds=API_KEY_HASH_ROUNDS)
+
+
+def build_displayable_api_key(api_key: str) -> str:
+    if api_key.startswith(_API_KEY_PREFIX):
+        api_key = api_key[len(_API_KEY_PREFIX) :]
+
+    return _API_KEY_PREFIX + api_key[:4] + "********" + api_key[-4:]
+
+
+def get_hashed_api_key_from_request(request: Request) -> str | None:
+    raw_api_key_header = request.headers.get(
+        _API_KEY_HEADER_ALTERNATIVE_NAME
+    ) or request.headers.get(_API_KEY_HEADER_NAME)
+    if raw_api_key_header is None:
+        return None
+
+    if raw_api_key_header.startswith(_BEARER_PREFIX):
+        raw_api_key_header = raw_api_key_header[len(_BEARER_PREFIX) :].strip()
+
+    return hash_api_key(raw_api_key_header)

backend/danswer/auth/invited_users.py

@@ -9,6 +9,7 @@ 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/users.py

@@ -48,20 +48,20 @@ 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 import text
 from sqlalchemy.orm import attributes
 from sqlalchemy.orm import Session
 
+from danswer.auth.api_key import get_hashed_api_key_from_request
 from danswer.auth.invited_users import get_invited_users
 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 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 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
@@ -75,6 +75,7 @@ from danswer.configs.constants import AuthType
 from danswer.configs.constants import DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
 from danswer.configs.constants import DANSWER_API_KEY_PREFIX
 from danswer.configs.constants import UNNAMED_KEY_PLACEHOLDER
+from danswer.db.api_key import fetch_user_for_api_key
 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
@@ -83,21 +84,27 @@ 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_ee_implementation_or_noop
 from danswer.utils.variable_functionality import fetch_versioned_implementation
-from shared_configs.configs import current_tenant_id
+from shared_configs.configs import async_return_default_schema
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 
 logger = setup_logger()
 
 
+class BasicAuthenticationError(HTTPException):
+    def __init__(self, detail: str):
+        super().__init__(status_code=status.HTTP_403_FORBIDDEN, detail=detail)
+
+
 def is_user_admin(user: User | None) -> bool:
     if AUTH_TYPE == AuthType.DISABLED:
         return True
@@ -132,7 +139,9 @@ 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 AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
+    return not DISABLE_VERIFICATION and (
+        AUTH_TYPE != AuthType.BASIC or REQUIRE_EMAIL_VERIFICATION
+    )
 
 
 def verify_email_is_invited(email: str) -> None:
@@ -185,20 +194,6 @@ def verify_email_domain(email: str) -> None:
             )
 
 
-def get_tenant_id_for_email(email: str) -> str:
-    if not MULTI_TENANT:
-        return "public"
-    # 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,
@@ -233,20 +228,16 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         safe: bool = False,
         request: Optional[Request] = None,
     ) -> User:
-        try:
-            tenant_id = (
-                get_tenant_id_for_email(user_create.email) if MULTI_TENANT else "public"
-            )
-        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"
-            )
+        tenant_id = await fetch_ee_implementation_or_noop(
+            "danswer.server.tenants.provisioning",
+            "get_or_create_tenant_id",
+            async_return_default_schema,
+        )(
+            email=user_create.email,
+        )
 
         async with get_async_session_with_tenant(tenant_id) as db_session:
-            token = current_tenant_id.set(tenant_id)
+            token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
 
             verify_email_is_invited(user_create.email)
             verify_email_domain(user_create.email)
@@ -264,7 +255,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                     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:
@@ -285,32 +276,11 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                 else:
                     raise exceptions.UserAlreadyExists()
 
-            current_tenant_id.reset(token)
+            finally:
+                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",
-        )
-
     async def oauth_callback(
         self: "BaseUserManager[models.UOAP, models.ID]",
         oauth_name: str,
@@ -324,23 +294,25 @@ 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 "public"
-            )
-        except exceptions.UserNotExists:
-            raise HTTPException(status_code=401, detail="User not found")
+        tenant_id = await fetch_ee_implementation_or_noop(
+            "danswer.server.tenants.provisioning",
+            "get_or_create_tenant_id",
+            async_return_default_schema,
+        )(
+            email=account_email,
+        )
 
         if not tenant_id:
             raise HTTPException(status_code=401, detail="User not found")
 
+        # Proceed with the tenant context
         token = None
         async with get_async_session_with_tenant(tenant_id) as db_session:
-            token = current_tenant_id.set(tenant_id)
+            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
@@ -380,9 +352,13 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                     }
 
                     user = await self.user_db.create(user_dict)
-                    user = await self.user_db.add_oauth_account(
-                        user, oauth_account_dict
-                    )
+
+                    # Explicitly set the Postgres schema for this session to ensure
+                    # OAuth account creation happens in the correct tenant schema
+                    await db_session.execute(text(f'SET search_path = "{tenant_id}"'))
+
+                    # Add OAuth account
+                    await self.user_db.add_oauth_account(user, oauth_account_dict)
                     await self.on_after_register(user, request)
 
             else:
@@ -426,7 +402,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                 user.oidc_expiry = None  # type: ignore
 
             if token:
-                current_tenant_id.reset(token)
+                CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
             return user
 
@@ -462,7 +438,13 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         email = credentials.username
 
         # Get tenant_id from mapping table
-        tenant_id = get_tenant_id_for_email(email)
+        tenant_id = await fetch_ee_implementation_or_noop(
+            "danswer.server.tenants.provisioning",
+            "get_or_create_tenant_id",
+            async_return_default_schema,
+        )(
+            email=email,
+        )
         if not tenant_id:
             # User not found in mapping
             self.password_helper.hash(credentials.password)
@@ -486,8 +468,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             has_web_login = attributes.get_attribute(user, "has_web_login")
 
             if not has_web_login:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
+                raise BasicAuthenticationError(
                     detail="NO_WEB_LOGIN_AND_HAS_NO_PASSWORD",
                 )
 
@@ -517,8 +498,33 @@ cookie_transport = CookieTransport(
 )
 
 
-def get_jwt_strategy() -> JWTStrategy:
-    return JWTStrategy(
+# This strategy is used to add tenant_id to the JWT token
+class TenantAwareJWTStrategy(JWTStrategy):
+    async def _create_token_data(self, user: User, impersonate: bool = False) -> dict:
+        tenant_id = await fetch_ee_implementation_or_noop(
+            "danswer.server.tenants.provisioning",
+            "get_or_create_tenant_id",
+            async_return_default_schema,
+        )(
+            email=user.email,
+        )
+
+        data = {
+            "sub": str(user.id),
+            "aud": self.token_audience,
+            "tenant_id": tenant_id,
+        }
+        return data
+
+    async def write_token(self, user: User) -> str:
+        data = await self._create_token_data(user)
+        return generate_jwt(
+            data, self.encode_key, self.lifetime_seconds, algorithm=self.algorithm
+        )
+
+
+def get_jwt_strategy() -> TenantAwareJWTStrategy:
+    return TenantAwareJWTStrategy(
         secret=USER_AUTH_SECRET,
         lifetime_seconds=SESSION_EXPIRE_TIME_SECONDS,
     )
@@ -619,14 +625,12 @@ async def double_check_user(
         return None
 
     if user is None:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not authenticated.",
         )
 
     if user_needs_to_be_verified() and not user.is_verified:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not verified.",
         )
 
@@ -635,8 +639,7 @@ async def double_check_user(
         and user.oidc_expiry < datetime.now(timezone.utc)
         and not include_expired
     ):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User's OIDC token has expired.",
         )
 
@@ -662,15 +665,13 @@ async def current_curator_or_admin_user(
         return None
 
     if not user or not hasattr(user, "role"):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not authenticated or lacks role information.",
         )
 
     allowed_roles = {UserRole.GLOBAL_CURATOR, UserRole.CURATOR, UserRole.ADMIN}
     if user.role not in allowed_roles:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not a curator or admin.",
         )
 
@@ -682,8 +683,7 @@ async def current_admin_user(user: User | None = Depends(current_user)) -> User
         return None
 
     if not user or not hasattr(user, "role") or user.role != UserRole.ADMIN:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User must be an admin to perform this action.",
         )
 
@@ -876,3 +876,22 @@ def get_oauth_router(
         return redirect_response
 
     return router
+
+
+def api_key_dep(
+    request: Request, db_session: Session = Depends(get_session)
+) -> User | None:
+    if AUTH_TYPE == AuthType.DISABLED:
+        return None
+
+    hashed_api_key = get_hashed_api_key_from_request(request)
+    if not hashed_api_key:
+        raise HTTPException(status_code=401, detail="Missing API key")
+
+    if hashed_api_key:
+        user = fetch_user_for_api_key(hashed_api_key, db_session)
+
+    if user is None:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    return user

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

@@ -0,0 +1,384 @@
+import logging
+import multiprocessing
+import time
+from typing import Any
+
+import requests
+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 sqlalchemy import text
+from sqlalchemy.orm import Session
+
+from danswer.background.celery.apps.task_formatters import CeleryTaskColoredFormatter
+from danswer.background.celery.apps.task_formatters import CeleryTaskPlainFormatter
+from danswer.background.celery.celery_utils import celery_is_worker_primary
+from danswer.configs.constants import DanswerRedisLocks
+from danswer.db.engine import get_sqlalchemy_engine
+from danswer.document_index.vespa_constants import VESPA_CONFIG_SERVER_URL
+from danswer.redis.redis_connector import RedisConnector
+from danswer.redis.redis_connector_credential_pair import RedisConnectorCredentialPair
+from danswer.redis.redis_connector_delete import RedisConnectorDelete
+from danswer.redis.redis_connector_prune import RedisConnectorPrune
+from danswer.redis.redis_document_set import RedisDocumentSet
+from danswer.redis.redis_pool import get_redis_client
+from danswer.redis.redis_usergroup import RedisUserGroup
+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(tenant_id, 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(tenant_id, int(usergroup_id))
+            r.srem(rug.taskset_key, task_id)
+        return
+
+    if task_id.startswith(RedisConnectorDelete.PREFIX):
+        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            RedisConnectorDelete.remove_from_taskset(int(cc_pair_id), task_id, r)
+        return
+
+    if task_id.startswith(RedisConnectorPrune.SUBTASK_PREFIX):
+        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
+        if cc_pair_id is not None:
+            RedisConnectorPrune.remove_from_taskset(int(cc_pair_id), task_id, r)
+        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:
+    """Waits for redis to become ready subject to a hardcoded timeout.
+    Will raise WorkerShutdown to kill the celery worker if the timeout is reached."""
+
+    r = get_redis_client(tenant_id=None)
+
+    WAIT_INTERVAL = 5
+    WAIT_LIMIT = 60
+
+    ready = False
+    time_start = time.monotonic()
+    logger.info("Redis: Readiness probe starting.")
+    while True:
+        try:
+            if r.ping():
+                ready = True
+                break
+        except Exception:
+            pass
+
+        time_elapsed = time.monotonic() - time_start
+        if time_elapsed > WAIT_LIMIT:
+            break
+
+        logger.info(
+            f"Redis: Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
+        )
+
+        time.sleep(WAIT_INTERVAL)
+
+    if not ready:
+        msg = (
+            f"Redis: Readiness probe did not succeed within the timeout "
+            f"({WAIT_LIMIT} seconds). Exiting..."
+        )
+        logger.error(msg)
+        raise WorkerShutdown(msg)
+
+    logger.info("Redis: Readiness probe succeeded. Continuing...")
+    return
+
+
+def wait_for_db(sender: Any, **kwargs: Any) -> None:
+    """Waits for the db to become ready subject to a hardcoded timeout.
+    Will raise WorkerShutdown to kill the celery worker if the timeout is reached."""
+
+    WAIT_INTERVAL = 5
+    WAIT_LIMIT = 60
+
+    ready = False
+    time_start = time.monotonic()
+    logger.info("Database: Readiness probe starting.")
+    while True:
+        try:
+            with Session(get_sqlalchemy_engine()) as db_session:
+                result = db_session.execute(text("SELECT NOW()")).scalar()
+                if result:
+                    ready = True
+                    break
+        except Exception:
+            pass
+
+        time_elapsed = time.monotonic() - time_start
+        if time_elapsed > WAIT_LIMIT:
+            break
+
+        logger.info(
+            f"Database: Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
+        )
+
+        time.sleep(WAIT_INTERVAL)
+
+    if not ready:
+        msg = (
+            f"Database: Readiness probe did not succeed within the timeout "
+            f"({WAIT_LIMIT} seconds). Exiting..."
+        )
+        logger.error(msg)
+        raise WorkerShutdown(msg)
+
+    logger.info("Database: Readiness probe succeeded. Continuing...")
+    return
+
+
+def wait_for_vespa(sender: Any, **kwargs: Any) -> None:
+    """Waits for Vespa to become ready subject to a hardcoded timeout.
+    Will raise WorkerShutdown to kill the celery worker if the timeout is reached."""
+
+    WAIT_INTERVAL = 5
+    WAIT_LIMIT = 60
+
+    ready = False
+    time_start = time.monotonic()
+    logger.info("Vespa: Readiness probe starting.")
+    while True:
+        try:
+            response = requests.get(f"{VESPA_CONFIG_SERVER_URL}/state/v1/health")
+            response.raise_for_status()
+
+            response_dict = response.json()
+            if response_dict["status"]["code"] == "up":
+                ready = True
+                break
+        except Exception:
+            pass
+
+        time_elapsed = time.monotonic() - time_start
+        if time_elapsed > WAIT_LIMIT:
+            break
+
+        logger.info(
+            f"Vespa: Readiness probe ongoing. elapsed={time_elapsed:.1f} timeout={WAIT_LIMIT:.1f}"
+        )
+
+        time.sleep(WAIT_INTERVAL)
+
+    if not ready:
+        msg = (
+            f"Vespa: Readiness probe did not succeed within the timeout "
+            f"({WAIT_LIMIT} seconds). Exiting..."
+        )
+        logger.error(msg)
+        raise WorkerShutdown(msg)
+
+    logger.info("Vespa: Readiness probe succeeded. Continuing...")
+    return
+
+
+def on_secondary_worker_init(sender: Any, **kwargs: Any) -> None:
+    logger.info("Running as a secondary celery worker.")
+
+    # Set up variables for waiting on primary worker
+    WAIT_INTERVAL = 5
+    WAIT_LIMIT = 60
+    r = get_redis_client(tenant_id=None)
+    time_start = time.monotonic()
+
+    logger.info("Waiting for primary worker to be ready...")
+    while True:
+        if r.exists(DanswerRedisLocks.PRIMARY_WORKER):
+            break
+
+        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
+
+
+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 sender.primary_worker_lock:
+        return
+
+    logger.info("Releasing primary worker lock.")
+    lock = sender.primary_worker_lock
+    try:
+        if lock.owned():
+            try:
+                lock.release()
+                sender.primary_worker_lock = None
+            except Exception as e:
+                logger.error(f"Failed to release primary worker lock: {e}")
+    except Exception as e:
+        logger.error(f"Failed to check if primary worker lock is owned: {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

@@ -0,0 +1,162 @@
+from datetime import timedelta
+from typing import Any
+
+from celery import Celery
+from celery import signals
+from celery.beat import PersistentScheduler  # type: ignore
+from celery.signals import beat_init
+
+import danswer.background.celery.apps.app_base as app_base
+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
+from danswer.utils.variable_functionality import fetch_versioned_implementation
+from shared_configs.configs import MULTI_TENANT
+
+logger = setup_logger(__name__)
+
+celery_app = Celery(__name__)
+celery_app.config_from_object("danswer.background.celery.configs.beat")
+
+
+class DynamicTenantScheduler(PersistentScheduler):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        logger.info("Initializing DynamicTenantScheduler")
+        super().__init__(*args, **kwargs)
+        self._reload_interval = timedelta(minutes=2)
+        self._last_reload = self.app.now() - self._reload_interval
+        # Let the parent class handle store initialization
+        self.setup_schedule()
+        self._update_tenant_tasks()
+        logger.info(f"Set reload interval to {self._reload_interval}")
+
+    def setup_schedule(self) -> None:
+        logger.info("Setting up initial schedule")
+        super().setup_schedule()
+        logger.info("Initial schedule setup complete")
+
+    def tick(self) -> float:
+        retval = super().tick()
+        now = self.app.now()
+        if (
+            self._last_reload is None
+            or (now - self._last_reload) > self._reload_interval
+        ):
+            logger.info("Reload interval reached, initiating tenant task update")
+            self._update_tenant_tasks()
+            self._last_reload = now
+            logger.info("Tenant task update completed, reset reload timer")
+        return retval
+
+    def _update_tenant_tasks(self) -> None:
+        logger.info("Starting tenant task update process")
+        try:
+            logger.info("Fetching all tenant IDs")
+            tenant_ids = get_all_tenant_ids()
+            logger.info(f"Found {len(tenant_ids)} tenants")
+
+            logger.info("Fetching tasks to schedule")
+            tasks_to_schedule = fetch_versioned_implementation(
+                "danswer.background.celery.tasks.beat_schedule", "get_tasks_to_schedule"
+            )
+
+            new_beat_schedule: dict[str, dict[str, Any]] = {}
+
+            current_schedule = self.schedule.items()
+
+            existing_tenants = set()
+            for task_name, _ in current_schedule:
+                if "-" in task_name:
+                    existing_tenants.add(task_name.split("-")[-1])
+            logger.info(f"Found {len(existing_tenants)} existing tenants in schedule")
+
+            for tenant_id in tenant_ids:
+                if tenant_id not in existing_tenants:
+                    logger.info(f"Processing new tenant: {tenant_id}")
+
+                for task in tasks_to_schedule():
+                    task_name = f"{task['name']}-{tenant_id}"
+                    logger.debug(f"Creating task configuration for {task_name}")
+                    new_task = {
+                        "task": task["task"],
+                        "schedule": task["schedule"],
+                        "kwargs": {"tenant_id": tenant_id},
+                    }
+                    if options := task.get("options"):
+                        logger.debug(f"Adding options to task {task_name}: {options}")
+                        new_task["options"] = options
+                    new_beat_schedule[task_name] = new_task
+
+            if self._should_update_schedule(current_schedule, new_beat_schedule):
+                logger.info(
+                    "Schedule update required",
+                    extra={
+                        "new_tasks": len(new_beat_schedule),
+                        "current_tasks": len(current_schedule),
+                    },
+                )
+
+                # Create schedule entries
+                entries = {}
+                for name, entry in new_beat_schedule.items():
+                    entries[name] = self.Entry(
+                        name=name,
+                        app=self.app,
+                        task=entry["task"],
+                        schedule=entry["schedule"],
+                        options=entry.get("options", {}),
+                        kwargs=entry.get("kwargs", {}),
+                    )
+
+                # Update the schedule using the scheduler's methods
+                self.schedule.clear()
+                self.schedule.update(entries)
+
+                # Ensure changes are persisted
+                self.sync()
+
+                logger.info("Schedule update completed successfully")
+            else:
+                logger.info("Schedule is up to date, no changes needed")
+
+        except (AttributeError, KeyError):
+            logger.exception("Failed to process task configuration")
+        except Exception:
+            logger.exception("Unexpected error updating tenant tasks")
+
+    def _should_update_schedule(
+        self, current_schedule: dict, new_schedule: dict
+    ) -> bool:
+        """Compare schedules to determine if an update is needed."""
+        logger.debug("Comparing current and new schedules")
+        current_tasks = set(name for name, _ in current_schedule)
+        new_tasks = set(new_schedule.keys())
+        needs_update = current_tasks != new_tasks
+        logger.debug(f"Schedule update needed: {needs_update}")
+        return needs_update
+
+
+@beat_init.connect
+def on_beat_init(sender: Any, **kwargs: Any) -> None:
+    logger.info("beat_init signal received.")
+
+    # Celery beat shouldn't touch the db at all. But just setting a low minimum here.
+    SqlEngine.set_app_name(POSTGRES_CELERY_BEAT_APP_NAME)
+    SqlEngine.init_engine(pool_size=2, max_overflow=0)
+
+    # Startup checks are not needed in multi-tenant case
+    if MULTI_TENANT:
+        return
+
+    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_app.conf.beat_scheduler = DynamicTenantScheduler

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

@@ -0,0 +1,95 @@
+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
+from shared_configs.configs import MULTI_TENANT
+
+
+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=4, max_overflow=12)
+
+    # Startup checks are not needed in multi-tenant case
+    if MULTI_TENANT:
+        return
+
+    app_base.wait_for_redis(sender, **kwargs)
+    app_base.wait_for_db(sender, **kwargs)
+    app_base.wait_for_vespa(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

@@ -0,0 +1,95 @@
+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
+from shared_configs.configs import MULTI_TENANT
+
+
+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)
+
+    # Startup checks are not needed in multi-tenant case
+    if MULTI_TENANT:
+        return
+
+    app_base.wait_for_redis(sender, **kwargs)
+    app_base.wait_for_db(sender, **kwargs)
+    app_base.wait_for_vespa(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

@@ -0,0 +1,96 @@
+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
+from shared_configs.configs import MULTI_TENANT
+
+
+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)
+    # Startup checks are not needed in multi-tenant case
+    if MULTI_TENANT:
+        return
+
+    app_base.wait_for_redis(sender, **kwargs)
+    app_base.wait_for_db(sender, **kwargs)
+    app_base.wait_for_vespa(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",
+        "danswer.background.celery.tasks.connector_deletion",
+    ]
+)

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

@@ -0,0 +1,240 @@
+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_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 SqlEngine
+from danswer.redis.redis_connector_credential_pair import RedisConnectorCredentialPair
+from danswer.redis.redis_connector_delete import RedisConnectorDelete
+from danswer.redis.redis_connector_index import RedisConnectorIndex
+from danswer.redis.redis_connector_prune import RedisConnectorPrune
+from danswer.redis.redis_connector_stop import RedisConnectorStop
+from danswer.redis.redis_document_set import RedisDocumentSet
+from danswer.redis.redis_pool import get_redis_client
+from danswer.redis.redis_usergroup import RedisUserGroup
+from danswer.utils.logger import setup_logger
+from shared_configs.configs import MULTI_TENANT
+
+
+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)
+
+    # Startup checks are not needed in multi-tenant case
+    if MULTI_TENANT:
+        return
+
+    app_base.wait_for_redis(sender, **kwargs)
+    app_base.wait_for_db(sender, **kwargs)
+    app_base.wait_for_vespa(sender, **kwargs)
+
+    logger.info("Running as the primary celery worker.")
+
+    # This is singleton work that should be done on startup exactly once
+    # by the primary worker. This is unnecessary in the multi tenant scenario
+    r = get_redis_client(tenant_id=None)
+
+    # 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_lock = 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())
+
+    RedisDocumentSet.reset_all(r)
+
+    RedisUserGroup.reset_all(r)
+
+    RedisConnectorDelete.reset_all(r)
+
+    RedisConnectorPrune.reset_all(r)
+
+    RedisConnectorIndex.reset_all(r)
+
+    RedisConnectorStop.reset_all(r)
+
+
+@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_lock"):
+                return
+
+            lock = worker.primary_worker_lock
+
+            r = get_redis_client(tenant_id=None)
+
+            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 worker restart or lock expiration."
+                )
+                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.")
+                    worker.primary_worker_lock = lock
+                else:
+                    task_logger.error("Primary worker lock: Acquire failed!")
+                    raise TimeoutError("Primary worker lock 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/scheduler.py

@@ -0,0 +1,96 @@
+from datetime import timedelta
+from typing import Any
+
+from celery.beat import PersistentScheduler  # type: ignore
+from celery.utils.log import get_task_logger
+
+from danswer.db.engine import get_all_tenant_ids
+from danswer.utils.variable_functionality import fetch_versioned_implementation
+
+logger = get_task_logger(__name__)
+
+
+class DynamicTenantScheduler(PersistentScheduler):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._reload_interval = timedelta(minutes=1)
+        self._last_reload = self.app.now() - self._reload_interval
+
+    def setup_schedule(self) -> None:
+        super().setup_schedule()
+
+    def tick(self) -> float:
+        retval = super().tick()
+        now = self.app.now()
+        if (
+            self._last_reload is None
+            or (now - self._last_reload) > self._reload_interval
+        ):
+            logger.info("Reloading schedule to check for new tenants...")
+            self._update_tenant_tasks()
+            self._last_reload = now
+        return retval
+
+    def _update_tenant_tasks(self) -> None:
+        logger.info("Checking for tenant task updates...")
+        try:
+            tenant_ids = get_all_tenant_ids()
+            tasks_to_schedule = fetch_versioned_implementation(
+                "danswer.background.celery.tasks.beat_schedule", "get_tasks_to_schedule"
+            )
+
+            new_beat_schedule: dict[str, dict[str, Any]] = {}
+
+            current_schedule = getattr(self, "_store", {"entries": {}}).get(
+                "entries", {}
+            )
+
+            existing_tenants = set()
+            for task_name in current_schedule.keys():
+                if "-" in task_name:
+                    existing_tenants.add(task_name.split("-")[-1])
+
+            for tenant_id in tenant_ids:
+                if tenant_id not in existing_tenants:
+                    logger.info(f"Found new tenant: {tenant_id}")
+
+                for task in tasks_to_schedule():
+                    task_name = f"{task['name']}-{tenant_id}"
+                    new_task = {
+                        "task": task["task"],
+                        "schedule": task["schedule"],
+                        "kwargs": {"tenant_id": tenant_id},
+                    }
+                    if options := task.get("options"):
+                        new_task["options"] = options
+                    new_beat_schedule[task_name] = new_task
+
+            if self._should_update_schedule(current_schedule, new_beat_schedule):
+                logger.info(
+                    "Updating schedule",
+                    extra={
+                        "new_tasks": len(new_beat_schedule),
+                        "current_tasks": len(current_schedule),
+                    },
+                )
+                if not hasattr(self, "_store"):
+                    self._store: dict[str, dict] = {"entries": {}}
+                self.update_from_dict(new_beat_schedule)
+                logger.info(f"New schedule: {new_beat_schedule}")
+
+                logger.info("Tenant tasks updated successfully")
+            else:
+                logger.debug("No schedule updates needed")
+
+        except (AttributeError, KeyError):
+            logger.exception("Failed to process task configuration")
+        except Exception:
+            logger.exception("Unexpected error updating tenant tasks")
+
+    def _should_update_schedule(
+        self, current_schedule: dict, new_schedule: dict
+    ) -> bool:
+        """Compare schedules to determine if an update is needed."""
+        current_tasks = set(current_schedule.keys())
+        new_tasks = set(new_schedule.keys())
+        return current_tasks != new_tasks

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

@@ -0,0 +1,26 @@
+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

@@ -1,619 +0,0 @@
-import logging
-import multiprocessing
-import time
-from datetime import timedelta
-from typing import Any
-
-import sentry_sdk
-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 celeryd_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 sentry_sdk.integrations.celery import CeleryIntegration
-
-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.celery_utils import celery_is_worker_primary
-from danswer.background.celery.celery_utils import get_all_tenant_ids
-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_INDEXING_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 get_session_with_tenant
-from danswer.db.engine import SqlEngine
-from danswer.db.search_settings import get_current_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 ColoredFormatter
-from danswer.utils.logger import PlainFormatter
-from danswer.utils.logger import setup_logger
-from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
-from shared_configs.configs import MODEL_SERVER_PORT
-from shared_configs.configs import SENTRY_DSN
-
-logger = setup_logger()
-
-# use this within celery tasks to get celery task specific logging
-task_logger = get_task_logger(__name__)
-
-if SENTRY_DSN:
-    sentry_sdk.init(
-        dsn=SENTRY_DSN,
-        integrations=[CeleryIntegration()],
-        traces_sample_rate=0.5,
-    )
-    logger.info("Sentry initialized")
-else:
-    logger.debug("Sentry DSN not provided, skipping Sentry initialization")
-
-
-celery_app = Celery(__name__)
-celery_app.config_from_object(
-    "danswer.background.celery.celeryconfig"
-)  # Load configuration from 'celeryconfig.py'
-
-
-@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,
-    tenant_id: str | None = None,
-    kwargs: dict | None = None,
-    **kwds: Any,
-) -> None:
-    pass
-
-
-@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[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
-
-    # 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 ''}"
-    )
-
-    if state not in READY_STATES:
-        return
-
-    if not task_id:
-        return
-
-    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
-
-
-@celeryd_init.connect
-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
-
-
-@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:
-    logger.info("worker_init signal received.")
-    logger.info(f"Multiprocessing start method: {multiprocessing.get_start_method()}")
-
-    # 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)
-    elif hostname.startswith("indexing"):
-        SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_APP_NAME)
-        SqlEngine.init_engine(pool_size=8, max_overflow=0)
-        tenant_ids = get_all_tenant_ids()
-
-        for tenant_id in tenant_ids:
-            # TODO: why is this necessary for the indexer to do?
-            with get_session_with_tenant(tenant_id) 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.")
-    else:
-        SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME)
-        SqlEngine.init_engine(pool_size=8, max_overflow=0)
-
-    if not hasattr(sender, "primary_worker_locks"):
-        sender.primary_worker_locks = {}
-
-    tenant_ids = get_all_tenant_ids()
-
-    if not celery_is_worker_primary(sender):
-        logger.info("Running as a secondary celery worker.")
-        for tenant_id in tenant_ids:
-            r = get_redis_client(tenant_id=tenant_id)
-            WAIT_INTERVAL = 5
-            WAIT_LIMIT = 60
-            time_start = time.monotonic()
-            logger.notice("Redis: Readiness check starting.")
-            while True:
-                # Log all the locks in Redis
-                all_locks = r.keys("*")
-                logger.notice(f"Current Redis locks: {all_locks}")
-                if r.exists(DanswerRedisLocks.PRIMARY_WORKER):
-                    break
-                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 = (
-                        "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("Wait for primary worker completed successfully. Continuing...")
-        return  # Exit the function for secondary workers
-
-    for tenant_id in tenant_ids:
-        r = get_redis_client(tenant_id=tenant_id)
-
-        WAIT_INTERVAL = 5
-        WAIT_LIMIT = 60
-
-        time_start = time.monotonic()
-        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(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!")
-
-        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)
-
-
-# @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:
-    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 hasattr(sender, "primary_worker_locks"):
-        return
-
-    logger.info("Releasing primary worker lock.")
-    for tenant_id, lock in sender.primary_worker_locks.items():
-        logger.info(f"Releasing primary worker lock for tenant {tenant_id}.")
-        if lock.owned():
-            lock.release()
-    sender.primary_worker_locks = {}
-
-
-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 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
-
-
-class HubPeriodicTask(bootsteps.StartStopStep):
-    """Regularly reacquires the primary worker locks for all tenants 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.
-    """
-
-    # Requires the Hub component
-    requires = {"celery.worker.components:Hub"}
-
-    def __init__(self, worker: Any, **kwargs: Any) -> None:
-        super().__init__(worker, **kwargs)
-        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 as e:
-            task_logger.error(f"Error in periodic task: {e}")
-
-    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",
-    ]
-)
-
-#####
-# 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 id in tenant_ids:
-    for task in tasks_to_schedule:
-        task_name = f"{task['name']}-{id}"  # Unique name for each scheduled task
-        beat_schedule[task_name] = {
-            "task": task["task"],
-            "schedule": task["schedule"],
-            "options": task["options"],
-            "kwargs": {"tenant_id": 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/celery_redis.py

@@ -1,542 +1,10 @@
 # These are helper objects for tracking the keys we need to write in redis
-import time
-from abc import ABC
-from abc import abstractmethod
 from typing import cast
-from uuid import uuid4
 
-import redis
-from celery import Celery
 from redis import Redis
-from sqlalchemy.orm import Session
 
-from danswer.background.celery.celeryconfig import CELERY_SEPARATOR
-from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
+from danswer.background.celery.configs.base import CELERY_SEPARATOR
 from danswer.configs.constants import DanswerCeleryPriority
-from danswer.configs.constants import DanswerCeleryQueues
-from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
-from danswer.db.document import construct_document_select_for_connector_credential_pair
-from danswer.db.document import (
-    construct_document_select_for_connector_credential_pair_by_needs_sync,
-)
-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):
-    PREFIX = "base"
-    FENCE_PREFIX = PREFIX + "_fence"
-    TASKSET_PREFIX = PREFIX + "_taskset"
-
-    def __init__(self, id: str):
-        self._id: str = id
-
-    @property
-    def task_id_prefix(self) -> str:
-        return f"{self.PREFIX}_{self._id}"
-
-    @property
-    def fence_key(self) -> str:
-        # example: documentset_fence_1
-        return f"{self.FENCE_PREFIX}_{self._id}"
-
-    @property
-    def taskset_key(self) -> str:
-        # example: documentset_taskset_1
-        return f"{self.TASKSET_PREFIX}_{self._id}"
-
-    @staticmethod
-    def get_id_from_fence_key(key: str) -> str | None:
-        """
-        Extracts the object ID from a fence key in the format `PREFIX_fence_X`.
-
-        Args:
-            key (str): The fence key string.
-
-        Returns:
-            Optional[int]: The extracted ID if the key is in the correct format, otherwise None.
-        """
-        parts = key.split("_")
-        if len(parts) != 3:
-            return None
-
-        object_id = parts[2]
-        return object_id
-
-    @staticmethod
-    def get_id_from_task_id(task_id: str) -> str | None:
-        """
-        Extracts the object ID from a task ID string.
-
-        This method assumes the task ID is formatted as `prefix_objectid_suffix`, where:
-        - `prefix` is an arbitrary string (e.g., the name of the task or entity),
-        - `objectid` is the ID you want to extract,
-        - `suffix` is another arbitrary string (e.g., a UUID).
-
-        Example:
-            If the input `task_id` is `documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc`,
-            this method will return the string `"1"`.
-
-        Args:
-            task_id (str): The task ID string from which to extract the object ID.
-
-        Returns:
-            str | None: The extracted object ID if the task ID is in the correct format, otherwise None.
-        """
-        # example: task_id=documentset_1_cbfdc96a-80ca-4312-a242-0bb68da3c1dc
-        parts = task_id.split("_")
-        if len(parts) != 3:
-            return None
-
-        object_id = parts[1]
-        return object_id
-
-    @abstractmethod
-    def generate_tasks(
-        self,
-        celery_app: Celery,
-        db_session: Session,
-        redis_client: Redis,
-        lock: redis.lock.Lock,
-        tenant_id: str | None,
-    ) -> int | None:
-        pass
-
-
-class RedisDocumentSet(RedisObjectHelper):
-    PREFIX = "documentset"
-    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)
-        for doc in db_session.scalars(stmt).yield_per(1):
-            current_time = time.monotonic()
-            if current_time - last_lock_time >= (
-                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
-            ):
-                lock.reacquire()
-                last_lock_time = current_time
-
-            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # we prefix the task id so it's easier to keep track of who created the task
-            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
-
-            # add to the set BEFORE creating the task.
-            redis_client.sadd(self.taskset_key, custom_task_id)
-
-            result = celery_app.send_task(
-                "vespa_metadata_sync_task",
-                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
-                queue=DanswerCeleryQueues.VESPA_METADATA_SYNC,
-                task_id=custom_task_id,
-                priority=DanswerCeleryPriority.LOW,
-            )
-
-            async_results.append(result)
-
-        return len(async_results)
-
-
-class RedisUserGroup(RedisObjectHelper):
-    PREFIX = "usergroup"
-    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",
-                "construct_document_select_by_usergroup",
-            )
-        except ModuleNotFoundError:
-            return 0
-
-        stmt = construct_document_select_by_usergroup(int(self._id))
-        for doc in db_session.scalars(stmt).yield_per(1):
-            current_time = time.monotonic()
-            if current_time - last_lock_time >= (
-                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
-            ):
-                lock.reacquire()
-                last_lock_time = current_time
-
-            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # we prefix the task id so it's easier to keep track of who created the task
-            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
-
-            # add to the set BEFORE creating the task.
-            redis_client.sadd(self.taskset_key, custom_task_id)
-
-            result = celery_app.send_task(
-                "vespa_metadata_sync_task",
-                kwargs=dict(document_id=doc.id, tenant_id=tenant_id),
-                queue=DanswerCeleryQueues.VESPA_METADATA_SYNC,
-                task_id=custom_task_id,
-                priority=DanswerCeleryPriority.LOW,
-            )
-
-            async_results.append(result)
-
-        return len(async_results)
-
-
-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
-    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
-
-    @classmethod
-    def get_taskset_key(cls) -> str:
-        return RedisConnectorCredentialPair.TASKSET_PREFIX
-
-    @property
-    def taskset_key(self) -> str:
-        """Notice that this is intentionally reusing the same taskset for all
-        connector syncs"""
-        # example: connector_taskset
-        return f"{self.TASKSET_PREFIX}"
-
-    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 = []
-        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
-        if not cc_pair:
-            return None
-
-        stmt = construct_document_select_for_connector_credential_pair_by_needs_sync(
-            cc_pair.connector_id, cc_pair.credential_id
-        )
-        for doc in db_session.scalars(stmt).yield_per(1):
-            current_time = time.monotonic()
-            if current_time - last_lock_time >= (
-                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
-            ):
-                lock.reacquire()
-                last_lock_time = current_time
-
-            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # we prefix the task id so it's easier to keep track of who created the task
-            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
-
-            # add to the tracking taskset in redis BEFORE creating the celery task.
-            # note that for the moment we are using a single taskset key, not differentiated by cc_pair id
-            redis_client.sadd(
-                RedisConnectorCredentialPair.get_taskset_key(), custom_task_id
-            )
-
-            # 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),
-                queue=DanswerCeleryQueues.VESPA_METADATA_SYNC,
-                task_id=custom_task_id,
-                priority=DanswerCeleryPriority.MEDIUM,
-            )
-
-            async_results.append(result)
-
-        return len(async_results)
-
-
-class RedisConnectorDeletion(RedisObjectHelper):
-    PREFIX = "connectordeletion"
-    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 = []
-        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
-        if not cc_pair:
-            return None
-
-        stmt = construct_document_select_for_connector_credential_pair(
-            cc_pair.connector_id, cc_pair.credential_id
-        )
-        for doc in db_session.scalars(stmt).yield_per(1):
-            current_time = time.monotonic()
-            if current_time - last_lock_time >= (
-                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
-            ):
-                lock.reacquire()
-                last_lock_time = current_time
-
-            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # the actual redis key is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # we prefix the task id so it's easier to keep track of who created the task
-            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            custom_task_id = f"{self.task_id_prefix}_{uuid4()}"
-
-            # add to the tracking taskset in redis BEFORE creating the celery task.
-            # note that for the moment we are using a single taskset key, not differentiated by cc_pair id
-            redis_client.sadd(self.taskset_key, custom_task_id)
-
-            # Priority on sync's triggered by new indexing should be medium
-            result = celery_app.send_task(
-                "document_by_cc_pair_cleanup_task",
-                kwargs=dict(
-                    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,
-                priority=DanswerCeleryPriority.MEDIUM,
-            )
-
-            async_results.append(result)
-
-        return len(async_results)
-
-
-class RedisConnectorPruning(RedisObjectHelper):
-    """Celery will kick off a long running generator task to crawl the connector and
-    find any missing docs, which will each then get a new cleanup task. The progress of
-    those tasks will then be monitored to completion.
-
-    Example rough happy path order:
-    Check connectorpruning_fence_1
-    Send generator task with id connectorpruning+generator_1_{uuid}
-
-    generator runs connector with callbacks that increment connectorpruning_generator_progress_1
-    generator creates many subtasks with id connectorpruning+sub_1_{uuid}
-      in taskset connectorpruning_taskset_1
-    on completion, generator sets connectorpruning_generator_complete_1
-
-    celery postrun removes subtasks from taskset
-    monitor beat task cleans up when taskset reaches 0 items
-    """
-
-    PREFIX = "connectorpruning"
-    FENCE_PREFIX = PREFIX + "_fence"  # a fence for the entire pruning process
-    GENERATOR_TASK_PREFIX = PREFIX + "+generator"
-
-    TASKSET_PREFIX = PREFIX + "_taskset"  # stores a list of prune tasks id's
-    SUBTASK_PREFIX = PREFIX + "+sub"
-
-    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, id: int) -> None:
-        super().__init__(str(id))
-        self.documents_to_prune: set[str] = set()
-
-    @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:
-        last_lock_time = time.monotonic()
-
-        async_results = []
-        cc_pair = get_connector_credential_pair_from_id(int(self._id), db_session)
-        if not cc_pair:
-            return None
-
-        for doc_id in self.documents_to_prune:
-            current_time = time.monotonic()
-            if lock and current_time - last_lock_time >= (
-                CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
-            ):
-                lock.reacquire()
-                last_lock_time = current_time
-
-            # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # the actual redis key is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            # we prefix the task id so it's easier to keep track of who created the task
-            # aka "documentset_1_6dd32ded3-00aa-4884-8b21-42f8332e7fac"
-            custom_task_id = f"{self.subtask_id_prefix}_{uuid4()}"
-
-            # add to the tracking taskset in redis BEFORE creating the celery task.
-            # note that for the moment we are using a single taskset key, not differentiated by cc_pair id
-            redis_client.sadd(self.taskset_key, custom_task_id)
-
-            # Priority on sync's triggered by new indexing should be medium
-            result = celery_app.send_task(
-                "document_by_cc_pair_cleanup_task",
-                kwargs=dict(
-                    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,
-                priority=DanswerCeleryPriority.MEDIUM,
-            )
-
-            async_results.append(result)
-
-        return len(async_results)
-
-    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=int(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 celery_get_queue_length(queue: str, r: Redis) -> int:

backend/danswer/background/celery/celery_utils.py

@@ -1,28 +1,23 @@
-from collections.abc import Callable
 from datetime import datetime
 from datetime import timezone
 from typing import Any
 
-from sqlalchemy import text
 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.configs.app_configs import MULTI_TENANT
-from danswer.configs.constants import TENANT_ID_PREFIX
 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.engine import get_session_with_tenant
 from danswer.db.enums import TaskStatus
 from danswer.db.models import TaskQueueState
-from danswer.redis.redis_pool import get_redis_client
+from danswer.redis.redis_connector import RedisConnector
 from danswer.server.documents.models import DeletionAttemptSnapshot
 from danswer.utils.logger import setup_logger
 
@@ -45,14 +40,14 @@ def _get_deletion_status(
     if not cc_pair:
         return None
 
-    rcd = RedisConnectorDeletion(cc_pair.id)
-
-    r = get_redis_client(tenant_id=tenant_id)
-    if not r.exists(rcd.fence_key):
+    redis_connector = RedisConnector(tenant_id, cc_pair.id)
+    if not redis_connector.delete.fenced:
         return None
 
     return TaskQueueState(
-        task_id="", task_name=rcd.fence_key, status=TaskStatus.STARTED
+        task_id="",
+        task_name=redis_connector.delete.fence_key,
+        status=TaskStatus.STARTED,
     )
 
 
@@ -75,13 +70,15 @@ 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,
-    progress_callback: Callable[[int], None] | None = None,
+    callback: RunIndexingCallbackInterface | None = None,
 ) -> set[str]:
     """
     If the PruneConnector hasnt been implemented for the given connector, just pull
@@ -91,10 +88,13 @@ 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, IdConnector):
-        all_connector_doc_ids = runnable_connector.retrieve_all_source_ids()
-    elif isinstance(runnable_connector, LoadConnector):
+
+    if 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()
@@ -103,16 +103,17 @@ def extract_ids_from_runnable_connector(
     else:
         raise RuntimeError("Pruning job could not find a valid runnable_connector.")
 
-    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))
+    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))
 
     return all_connector_doc_ids
 
@@ -133,33 +134,10 @@ 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 either in celeryconfig.py or on the
+    for the celery worker, which can be done on the
     command line with '--hostname'."""
     hostname = worker.hostname
     if hostname.startswith("primary"):
         return True
 
     return False
-
-
-def get_all_tenant_ids() -> list[str] | list[None]:
-    if not MULTI_TENANT:
-        return [None]
-    with get_session_with_tenant(tenant_id="public") as session:
-        result = session.execute(
-            text(
-                """
-            SELECT schema_name
-            FROM information_schema.schemata
-            WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'public')"""
-            )
-        )
-        tenant_ids = [row[0] for row in result]
-
-    valid_tenants = [
-        tenant
-        for tenant in tenant_ids
-        if tenant is None or tenant.startswith(TENANT_ID_PREFIX)
-    ]
-
-    return valid_tenants

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

@@ -31,21 +31,10 @@ 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
-
-# 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
-
 broker_connection_retry_on_startup = True
 broker_pool_limit = CELERY_BROKER_POOL_LIMIT
 
@@ -60,6 +49,7 @@ 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
@@ -73,10 +63,19 @@ 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
 
@@ -102,3 +101,4 @@ 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

@@ -0,0 +1,14 @@
+# 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

@@ -0,0 +1,20 @@
+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

@@ -0,0 +1,21 @@
+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

@@ -0,0 +1,22 @@
+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

@@ -0,0 +1,20 @@
+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/beat_schedule.py

@@ -0,0 +1,48 @@
+from datetime import timedelta
+from typing import Any
+
+from danswer.configs.constants import DanswerCeleryPriority
+
+
+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=20),
+        "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},
+    },
+]
+
+
+def get_tasks_to_schedule() -> list[dict[str, Any]]:
+    return tasks_to_schedule

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

@@ -1,29 +1,40 @@
+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.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.apps.app_base import task_logger
 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.enums import ConnectorCredentialPairStatus
-from danswer.db.models import ConnectorCredentialPair
+from danswer.db.search_settings import get_all_search_settings
+from danswer.redis.redis_connector import RedisConnector
+from danswer.redis.redis_connector_delete import RedisConnectorDeletionFenceData
 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(*, tenant_id: str | None) -> None:
+def check_for_connector_deletion_task(self: Task, *, tenant_id: str | None) -> None:
     r = get_redis_client(tenant_id=tenant_id)
 
     lock_beat = r.lock(
@@ -36,25 +47,44 @@ def check_for_connector_deletion_task(*, tenant_id: str | None) -> None:
         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:
             cc_pairs = get_connector_credential_pairs(db_session)
             for cc_pair in cc_pairs:
-                try_generate_document_cc_pair_cleanup_tasks(
-                    cc_pair, db_session, r, lock_beat, tenant_id
-                )
+                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:
+                redis_connector = RedisConnector(tenant_id, 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))
+                    redis_connector.stop.set_fence(True)
+                else:
+                    # clear the stop signal if it exists ... no longer needed
+                    redis_connector.stop.set_fence(False)
+
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
         )
     except Exception:
-        task_logger.exception("Unexpected exception")
+        task_logger.exception(f"Unexpected exception: tenant={tenant_id}")
     finally:
         if lock_beat.owned():
             lock_beat.release()
 
 
 def try_generate_document_cc_pair_cleanup_tasks(
-    cc_pair: ConnectorCredentialPair,
+    app: Celery,
+    cc_pair_id: int,
     db_session: Session,
     r: Redis,
     lock_beat: redis.lock.Lock,
@@ -63,51 +93,89 @@ def try_generate_document_cc_pair_cleanup_tasks(
     """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)
+    redis_connector = RedisConnector(tenant_id, cc_pair_id)
 
     # don't generate sync tasks if tasks are still pending
-    if r.exists(rcd.fence_key):
+    if redis_connector.delete.fenced:
         return None
 
-    # we need to refresh the state of the object inside the fence
+    # we need to load the state of the object inside the fence
     # to avoid a race condition with db.commit/fence deletion
     # at the end of this taskset
-    try:
-        db_session.refresh(cc_pair)
-    except ObjectDeletedError:
+    cc_pair = get_connector_credential_pair_from_id(cc_pair_id, db_session)
+    if not cc_pair:
         return None
 
     if cc_pair.status != ConnectorCredentialPairStatus.DELETING:
         return None
 
-    # add tasks to celery and build up the task set to monitor in redis
-    r.delete(rcd.taskset_key)
+    # set a basic fence to start
+    fence_payload = RedisConnectorDeletionFenceData(
+        num_tasks=None,
+        submitted=datetime.now(timezone.utc),
+    )
 
-    # 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}"
-    )
-    tasks_generated = rcd.generate_tasks(
-        celery_app, db_session, r, lock_beat, tenant_id
-    )
-    if tasks_generated is None:
+    redis_connector.delete.set_fence(fence_payload)
+
+    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:
+            redis_connector_index = redis_connector.new_index(search_settings.id)
+            if redis_connector_index.fenced:
+                raise TaskDependencyError(
+                    f"Connector deletion - Delayed (indexing in progress): "
+                    f"cc_pair={cc_pair_id} "
+                    f"search_settings={search_settings.id}"
+                )
+
+        if redis_connector.prune.fenced:
+            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
+        redis_connector.delete.taskset_clear()
+
+        # 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 = redis_connector.delete.generate_tasks(
+            app, db_session, lock_beat
+        )
+        if tasks_generated is None:
+            raise ValueError("RedisConnectorDeletion.generate_tasks returned None")
+    except TaskDependencyError:
+        redis_connector.delete.set_fence(None)
+        raise
+    except Exception:
+        task_logger.exception("Unexpected exception")
+        redis_connector.delete.set_fence(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
 
-    # 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}"
+        )
 
-    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
+        fence_payload.num_tasks = tasks_generated
+        redis_connector.delete.set_fence(fence_payload)
 
-    # 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

@@ -2,20 +2,20 @@ 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
+import sentry_sdk
+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.celery_app import celery_app
-from danswer.background.celery.celery_app import task_logger
-from danswer.background.celery.celery_redis import RedisConnectorIndexing
-from danswer.background.celery.tasks.shared.tasks import RedisConnectorIndexingFenceData
+from danswer.background.celery.apps.app_base import task_logger
 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.constants import CELERY_INDEXING_LOCK_TIMEOUT
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
@@ -40,18 +40,52 @@ 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_connector import RedisConnector
+from danswer.redis.redis_connector_index import RedisConnectorIndexingFenceData
 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
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.configs import SENTRY_DSN
 
 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(*, tenant_id: str | None) -> int | None:
+def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
     tasks_created = 0
 
     r = get_redis_client(tenant_id=tenant_id)
@@ -64,29 +98,56 @@ def check_for_indexing(*, tenant_id: str | None) -> int | None:
     try:
         # these tasks should never overlap
         if not lock_beat.acquire(blocking=False):
-            task_logger.info(f"Lock acquired for tenant (Y): {tenant_id}")
             return None
-        else:
-            task_logger.info(f"Lock acquired for tenant (N): {tenant_id}")
 
-        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)
-
-            cc_pairs = fetch_connector_credential_pairs(db_session)
-            for cc_pair in cc_pairs:
-                for search_settings_instance in search_settings:
-                    rci = RedisConnectorIndexing(
-                        cc_pair.id, search_settings_instance.id
+        with get_session_with_tenant(tenant_id=tenant_id) as db_session:
+            old_search_settings = 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:
+                if old_search_settings:
+                    embedding_model = EmbeddingModel.from_db_model(
+                        search_settings=current_search_settings,
+                        server_host=INDEXING_MODEL_SERVER_HOST,
+                        server_port=INDEXING_MODEL_SERVER_PORT,
                     )
-                    if r.exists(rci.fence_key):
+
+                    # only warm up if search settings were changed
+                    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:
+            redis_connector = RedisConnector(tenant_id, cc_pair_id)
+            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:
+                    redis_connector_index = redis_connector.new_index(
+                        search_settings_instance.id
+                    )
+                    if redis_connector_index.fenced:
+                        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(
@@ -104,6 +165,7 @@ def check_for_indexing(*, tenant_id: str | None) -> int | None:
                     # 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,
@@ -113,7 +175,9 @@ def check_for_indexing(*, tenant_id: str | None) -> int | None:
                     )
                     if attempt_id:
                         task_logger.info(
-                            f"Indexing queued: cc_pair_id={cc_pair.id} index_attempt_id={attempt_id}"
+                            f"Indexing queued: index_attempt={attempt_id} "
+                            f"cc_pair={cc_pair.id} "
+                            f"search_settings={search_settings_instance.id} "
                         )
                         tasks_created += 1
     except SoftTimeLimitExceeded:
@@ -121,7 +185,7 @@ def check_for_indexing(*, tenant_id: str | None) -> int | None:
             "Soft time limit exceeded, task is being terminated gracefully."
         )
     except Exception:
-        task_logger.exception("Unexpected exception")
+        task_logger.exception(f"Unexpected exception: tenant={tenant_id}")
     finally:
         if lock_beat.owned():
             lock_beat.release()
@@ -213,6 +277,7 @@ def _should_index(
 
 
 def try_creating_indexing_task(
+    celery_app: Celery,
     cc_pair: ConnectorCredentialPair,
     search_settings: SearchSettings,
     reindex: bool,
@@ -241,24 +306,38 @@ def try_creating_indexing_task(
         return None
 
     try:
-        rci = RedisConnectorIndexing(cc_pair.id, search_settings.id)
+        redis_connector = RedisConnector(tenant_id, cc_pair.id)
+        redis_connector_index = redis_connector.new_index(search_settings.id)
 
         # skip if already indexing
-        if r.exists(rci.fence_key):
+        if redis_connector_index.fenced:
             return None
 
         # skip indexing if the cc_pair is deleting
+        if redis_connector.delete.fenced:
+            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)
+        redis_connector_index.generator_clear()
 
-        custom_task_id = f"{rci.generator_task_id_prefix}_{uuid4()}"
+        # set a basic fence to start
+        payload = RedisConnectorIndexingFenceData(
+            index_attempt_id=None,
+            started=None,
+            submitted=datetime.now(timezone.utc),
+            celery_task_id=None,
+        )
 
-        # create the index attempt ... just for tracking purposes
+        redis_connector_index.set_fence(payload)
+
+        # 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,
@@ -266,6 +345,8 @@ def try_creating_indexing_task(
             db_session=db_session,
         )
 
+        custom_task_id = redis_connector_index.generate_generator_task_id()
+
         result = celery_app.send_task(
             "connector_indexing_proxy_task",
             kwargs=dict(
@@ -279,18 +360,21 @@ def try_creating_indexing_task(
             priority=DanswerCeleryPriority.MEDIUM,
         )
         if not result:
-            return None
+            raise RuntimeError("send_task for connector_indexing_proxy_task failed.")
+
+        # now fill out the fence with the rest of the data
+        payload.index_attempt_id = index_attempt_id
+        payload.celery_task_id = result.id
+        redis_connector_index.set_fence(payload)
 
-        # set this only after all tasks have been added
-        fence_value = RedisConnectorIndexingFenceData(
-            index_attempt_id=index_attempt_id,
-            started=None,
-            submitted=datetime.now(timezone.utc),
-            celery_task_id=result.id,
-        )
-        r.set(rci.fence_key, fence_value.model_dump_json())
     except Exception:
-        task_logger.exception("Unexpected exception")
+        redis_connector_index.set_fence(payload)
+        task_logger.exception(
+            f"Unexpected exception: "
+            f"tenant={tenant_id} "
+            f"cc_pair={cc_pair.id} "
+            f"search_settings={search_settings.id}"
+        )
         return None
     finally:
         if lock.owned():
@@ -307,7 +391,12 @@ def connector_indexing_proxy_task(
     tenant_id: str | None,
 ) -> None:
     """celery tasks are forked, but forking is unstable.  This proxies work to a spawned task."""
-
+    task_logger.info(
+        f"Indexing proxy - starting: attempt={index_attempt_id} "
+        f"tenant={tenant_id} "
+        f"cc_pair={cc_pair_id} "
+        f"search_settings={search_settings_id}"
+    )
     client = SimpleJobClient()
 
     job = client.submit(
@@ -321,29 +410,56 @@ def connector_indexing_proxy_task(
     )
 
     if not job:
+        task_logger.info(
+            f"Indexing proxy - spawn failed: attempt={index_attempt_id} "
+            f"tenant={tenant_id} "
+            f"cc_pair={cc_pair_id} "
+            f"search_settings={search_settings_id}"
+        )
         return
 
+    task_logger.info(
+        f"Indexing proxy - spawn succeeded: attempt={index_attempt_id} "
+        f"tenant={tenant_id} "
+        f"cc_pair={cc_pair_id} "
+        f"search_settings={search_settings_id}"
+    )
+
     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():
+        # do nothing for ongoing jobs that haven't been stopped
+        if not job.done():
+            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
+                )
+
                 if not index_attempt:
                     continue
 
                 if not index_attempt.is_finished():
                     continue
 
-            if job.status == "error":
-                logger.error(job.exception())
+        if job.status == "error":
+            task_logger.error(
+                f"Indexing proxy - spawned task exceptioned: "
+                f"attempt={index_attempt_id} "
+                f"tenant={tenant_id} "
+                f"cc_pair={cc_pair_id} "
+                f"search_settings={search_settings_id} "
+                f"error={job.exception()}"
+            )
 
-            job.release()
-            break
+        job.release()
+        break
 
+    task_logger.info(
+        f"Indexing proxy - finished: attempt={index_attempt_id} "
+        f"tenant={tenant_id} "
+        f"cc_pair={cc_pair_id} "
+        f"search_settings={search_settings_id}"
+    )
     return
 
 
@@ -365,36 +481,106 @@ def connector_indexing_task(
 
     Returns None if the task did not run (possibly due to a conflict).
     Otherwise, returns an int >= 0 representing the number of indexed docs.
+
+    NOTE: if an exception is raised out of this task, the primary worker will detect
+    that the task transitioned to a "READY" state but the generator_complete_key doesn't exist.
+    This will cause the primary worker to abort the indexing attempt and clean up.
     """
 
-    attempt = None
-    n_final_progress = 0
+    # Since connector_indexing_proxy_task spawns a new process using this function as
+    # the entrypoint, we init Sentry here.
+    if SENTRY_DSN:
+        sentry_sdk.init(
+            dsn=SENTRY_DSN,
+            traces_sample_rate=0.1,
+        )
+        logger.info("Sentry initialized")
+    else:
+        logger.debug("Sentry DSN not provided, skipping Sentry initialization")
+
+    logger.info(
+        f"Indexing spawned task starting: attempt={index_attempt_id} "
+        f"tenant={tenant_id} "
+        f"cc_pair={cc_pair_id} "
+        f"search_settings={search_settings_id}"
+    )
+
+    attempt_found = False
+    n_final_progress: int | None = None
+
+    redis_connector = RedisConnector(tenant_id, cc_pair_id)
+    redis_connector_index = redis_connector.new_index(search_settings_id)
 
     r = get_redis_client(tenant_id=tenant_id)
 
-    rci = RedisConnectorIndexing(cc_pair_id, search_settings_id)
+    if redis_connector.delete.fenced:
+        raise RuntimeError(
+            f"Indexing will not start because connector deletion is in progress: "
+            f"cc_pair={cc_pair_id} "
+            f"fence={redis_connector.delete.fence_key}"
+        )
+
+    if redis_connector.stop.fenced:
+        raise RuntimeError(
+            f"Indexing will not start because a connector stop signal was detected: "
+            f"cc_pair={cc_pair_id} "
+            f"fence={redis_connector.stop.fence_key}"
+        )
+
+    while True:
+        # wait for the fence to come up
+        if not redis_connector_index.fenced:
+            raise ValueError(
+                f"connector_indexing_task - fence not found: fence={redis_connector_index.fence_key}"
+            )
+
+        payload = redis_connector_index.payload
+        if not payload:
+            raise ValueError("connector_indexing_task: payload invalid or not found")
+
+        if payload.index_attempt_id is None or payload.celery_task_id is None:
+            logger.info(
+                f"connector_indexing_task - Waiting for fence: fence={redis_connector_index.fence_key}"
+            )
+            sleep(1)
+            continue
+
+        if payload.index_attempt_id != index_attempt_id:
+            raise ValueError(
+                f"connector_indexing_task - id mismatch. Task may be left over from previous run.: "
+                f"task_index_attempt={index_attempt_id} "
+                f"payload_index_attempt={payload.index_attempt_id}"
+            )
+
+        logger.info(
+            f"connector_indexing_task - Fence found, continuing...: fence={redis_connector_index.fence_key}"
+        )
+        break
 
     lock = r.lock(
-        rci.generator_lock_key,
+        redis_connector_index.generator_lock_key,
         timeout=CELERY_INDEXING_LOCK_TIMEOUT,
     )
 
     acquired = lock.acquire(blocking=False)
     if not acquired:
-        task_logger.warning(
+        logger.warning(
             f"Indexing task already running, exiting...: "
-            f"cc_pair_id={cc_pair_id} search_settings_id={search_settings_id}"
+            f"cc_pair={cc_pair_id} search_settings={search_settings_id}"
         )
-        # r.set(rci.generator_complete_key, HTTPStatus.CONFLICT.value)
         return None
 
+    payload.started = datetime.now(timezone.utc)
+    redis_connector_index.set_fence(payload)
+
     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_id={index_attempt_id}"
+                    f"Index attempt not found: index_attempt={index_attempt_id}"
                 )
+            attempt_found = True
 
             cc_pair = get_connector_credential_pair_from_id(
                 cc_pair_id=cc_pair_id,
@@ -402,54 +588,64 @@ def connector_indexing_task(
             )
 
             if not cc_pair:
-                raise ValueError(f"cc_pair not found: cc_pair_id={cc_pair_id}")
+                raise ValueError(f"cc_pair not found: cc_pair={cc_pair_id}")
 
             if not cc_pair.connector:
                 raise ValueError(
-                    f"Connector not found: connector_id={cc_pair.connector_id}"
+                    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: credential_id={cc_pair.credential_id}"
+                    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(
+            redis_connector.stop.fence_key,
+            redis_connector_index.generator_progress_key,
+            lock,
+            r,
+        )
 
-            # Define the callback function
-            def redis_increment_callback(amount: int) -> None:
-                lock.reacquire()
-                r.incrby(rci.generator_progress_key, amount)
+        logger.info(
+            f"Indexing spawned task running entrypoint: attempt={index_attempt_id} "
+            f"tenant={tenant_id} "
+            f"cc_pair={cc_pair_id} "
+            f"search_settings={search_settings_id}"
+        )
 
-            run_indexing_entrypoint(
-                index_attempt_id,
-                tenant_id,
-                cc_pair_id,
-                is_ee,
-                progress_callback=redis_increment_callback,
-            )
+        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)
+        # get back the total number of indexed docs and return it
+        n_final_progress = redis_connector_index.get_progress()
+        redis_connector_index.set_generator_complete(HTTPStatus.OK.value)
     except Exception as e:
-        task_logger.exception(f"Failed to run indexing for cc_pair_id={cc_pair_id}.")
-        if attempt:
-            mark_attempt_failed(attempt, db_session, failure_reason=str(e))
+        logger.exception(
+            f"Indexing spawned task failed: attempt={index_attempt_id} "
+            f"tenant={tenant_id} "
+            f"cc_pair={cc_pair_id} "
+            f"search_settings={search_settings_id}"
+        )
+        if attempt_found:
+            with get_session_with_tenant(tenant_id) as db_session:
+                mark_attempt_failed(index_attempt_id, 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()
 
+    logger.info(
+        f"Indexing spawned task finished: attempt={index_attempt_id} "
+        f"tenant={tenant_id} "
+        f"cc_pair={cc_pair_id} "
+        f"search_settings={search_settings_id}"
+    )
     return n_final_progress

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

@@ -11,7 +11,7 @@ from sqlalchemy import inspect
 from sqlalchemy import text
 from sqlalchemy.orm import Session
 
-from danswer.background.celery.celery_app import task_logger
+from danswer.background.celery.apps.app_base 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

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

@@ -3,15 +3,16 @@ from datetime import timedelta
 from datetime import timezone
 from uuid import uuid4
 
+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.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.apps.app_base import task_logger
 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
@@ -23,12 +24,15 @@ 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.enums import ConnectorCredentialPairStatus
 from danswer.db.models import ConnectorCredentialPair
+from danswer.redis.redis_connector import RedisConnector
 from danswer.redis.redis_pool import get_redis_client
+from danswer.utils.logger import pruning_ctx
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -37,8 +41,9 @@ logger = setup_logger()
 @shared_task(
     name="check_for_pruning",
     soft_time_limit=JOB_TIMEOUT,
+    bind=True,
 )
-def check_for_pruning(*, tenant_id: str | None) -> None:
+def check_for_pruning(self: Task, *, tenant_id: str | None) -> None:
     r = get_redis_client(tenant_id=tenant_id)
 
     lock_beat = r.lock(
@@ -51,26 +56,35 @@ def check_for_pruning(*, 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:
             cc_pairs = get_connector_credential_pairs(db_session)
-            for cc_pair in cc_pairs:
-                lock_beat.reacquire()
+            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(
-                    cc_pair, db_session, r, tenant_id
+                    self.app, cc_pair, db_session, r, tenant_id
                 )
                 if not tasks_created:
                     continue
 
-                task_logger.info(f"Pruning queued: cc_pair_id={cc_pair.id}")
+                task_logger.info(f"Pruning queued: cc_pair={cc_pair.id}")
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
         )
     except Exception:
-        task_logger.exception("Unexpected exception")
+        task_logger.exception(f"Unexpected exception: tenant={tenant_id}")
     finally:
         if lock_beat.owned():
             lock_beat.release()
@@ -118,6 +132,7 @@ def is_pruning_due(
 
 
 def try_creating_prune_generator_task(
+    celery_app: Celery,
     cc_pair: ConnectorCredentialPair,
     db_session: Session,
     r: Redis,
@@ -130,8 +145,11 @@ def try_creating_prune_generator_task(
     is used to trigger prunes immediately, e.g. via the web ui.
     """
 
+    redis_connector = RedisConnector(tenant_id, cc_pair.id)
+
     if not ALLOW_SIMULTANEOUS_PRUNING:
-        for key in r.scan_iter(RedisConnectorPruning.FENCE_PREFIX + "*"):
+        count = redis_connector.prune.get_active_task_count()
+        if count > 0:
             return None
 
     LOCK_TIMEOUT = 30
@@ -148,22 +166,21 @@ def try_creating_prune_generator_task(
         return None
 
     try:
-        rcp = RedisConnectorPruning(cc_pair.id)
-
-        # skip pruning if already pruning
-        if r.exists(rcp.fence_key):
+        if redis_connector.prune.fenced:  # skip pruning if already pruning
+            return None
+
+        if redis_connector.delete.fenced:  # skip pruning if the cc_pair is deleting
             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)
+        redis_connector.prune.generator_clear()
+        redis_connector.prune.taskset_clear()
 
-        custom_task_id = f"{rcp.generator_task_id_prefix}_{uuid4()}"
+        custom_task_id = f"{redis_connector.prune.generator_task_key}_{uuid4()}"
 
         celery_app.send_task(
             "connector_pruning_generator_task",
@@ -179,9 +196,9 @@ def try_creating_prune_generator_task(
         )
 
         # set this only after all tasks have been added
-        r.set(rcp.fence_key, 1)
+        redis_connector.prune.set_fence(True)
     except Exception:
-        task_logger.exception("Unexpected exception")
+        task_logger.exception(f"Unexpected exception: cc_pair={cc_pair.id}")
         return None
     finally:
         if lock.owned():
@@ -196,27 +213,37 @@ def try_creating_prune_generator_task(
     soft_time_limit=JOB_TIMEOUT,
     track_started=True,
     trail=False,
+    bind=True,
 )
 def connector_pruning_generator_task(
-    cc_pair_id: int, connector_id: int, credential_id: int, tenant_id: str | None
+    self: Task,
+    cc_pair_id: int,
+    connector_id: int,
+    credential_id: int,
+    tenant_id: str | None,
 ) -> 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"""
 
+    pruning_ctx_dict = pruning_ctx.get()
+    pruning_ctx_dict["cc_pair_id"] = cc_pair_id
+    pruning_ctx_dict["request_id"] = self.request.id
+    pruning_ctx.set(pruning_ctx_dict)
+
+    redis_connector = RedisConnector(tenant_id, cc_pair_id)
+
     r = get_redis_client(tenant_id=tenant_id)
 
-    rcp = RedisConnectorPruning(cc_pair_id)
-
     lock = r.lock(
-        DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{rcp._id}",
+        DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{redis_connector.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_id={cc_pair_id}"
+            f"Pruning task already running, exiting...: cc_pair={cc_pair_id}"
         )
         return None
 
@@ -234,22 +261,23 @@ def connector_pruning_generator_task(
                 )
                 return
 
-            # Define the callback function
-            def redis_increment_callback(amount: int) -> None:
-                lock.reacquire()
-                r.incrby(rcp.generator_progress_key, amount)
-
             runnable_connector = instantiate_connector(
                 db_session,
                 cc_pair.connector.source,
-                InputType.PRUNE,
+                InputType.SLIM_RETRIEVAL,
                 cc_pair.connector.connector_specific_config,
                 cc_pair.credential,
             )
 
+            callback = RunIndexingCallback(
+                redis_connector.stop.fence_key,
+                redis_connector.prune.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, redis_increment_callback
+                runnable_connector, callback
             )
 
             # a list of docs in our local index
@@ -267,34 +295,34 @@ def connector_pruning_generator_task(
 
             task_logger.info(
                 f"Pruning set collected: "
-                f"cc_pair_id={cc_pair.id} "
+                f"cc_pair={cc_pair_id} "
                 f"docs_to_remove={len(doc_ids_to_remove)} "
                 f"doc_source={cc_pair.connector.source}"
             )
 
-            rcp.documents_to_prune = set(doc_ids_to_remove)
-
             task_logger.info(
-                f"RedisConnectorPruning.generate_tasks starting. cc_pair_id={cc_pair.id}"
+                f"RedisConnector.prune.generate_tasks starting. cc_pair={cc_pair_id}"
             )
-            tasks_generated = rcp.generate_tasks(
-                celery_app, db_session, r, None, tenant_id
+            tasks_generated = redis_connector.prune.generate_tasks(
+                set(doc_ids_to_remove), self.app, db_session, None
             )
             if tasks_generated is None:
                 return None
 
             task_logger.info(
-                f"RedisConnectorPruning.generate_tasks finished. "
-                f"cc_pair_id={cc_pair.id} tasks_generated={tasks_generated}"
+                f"RedisConnector.prune.generate_tasks finished. "
+                f"cc_pair={cc_pair_id} tasks_generated={tasks_generated}"
             )
 
-            r.set(rcp.generator_complete_key, tasks_generated)
+            redis_connector.prune.generator_complete = tasks_generated
     except Exception as e:
-        task_logger.exception(f"Failed to run pruning for connector id {connector_id}.")
+        task_logger.exception(
+            f"Failed to run pruning: cc_pair={cc_pair_id} connector={connector_id}"
+        )
 
-        r.delete(rcp.generator_progress_key)
-        r.delete(rcp.taskset_key)
-        r.delete(rcp.fence_key)
+        redis_connector.prune.generator_clear()
+        redis_connector.prune.taskset_clear()
+        redis_connector.prune.set_fence(False)
         raise e
     finally:
         if lock.owned():

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

@@ -0,0 +1,40 @@
+import httpx
+from tenacity import retry
+from tenacity import retry_if_exception_type
+from tenacity import stop_after_delay
+from tenacity import wait_random_exponential
+
+from danswer.document_index.interfaces import DocumentIndex
+from danswer.document_index.interfaces import VespaDocumentFields
+
+
+class RetryDocumentIndex:
+    """A wrapper class to help with specific retries against Vespa involving
+    read timeouts.
+
+    wait_random_exponential implements full jitter as per this article:
+    https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/"""
+
+    MAX_WAIT = 30
+
+    # STOP_AFTER + MAX_WAIT should be slightly less (5?) than the celery soft_time_limit
+    STOP_AFTER = 70
+
+    def __init__(self, index: DocumentIndex):
+        self.index: DocumentIndex = index
+
+    @retry(
+        retry=retry_if_exception_type(httpx.ReadTimeout),
+        wait=wait_random_exponential(multiplier=1, max=MAX_WAIT),
+        stop=stop_after_delay(STOP_AFTER),
+    )
+    def delete_single(self, doc_id: str) -> int:
+        return self.index.delete_single(doc_id)
+
+    @retry(
+        retry=retry_if_exception_type(httpx.ReadTimeout),
+        wait=wait_random_exponential(multiplier=1, max=MAX_WAIT),
+        stop=stop_after_delay(STOP_AFTER),
+    )
+    def update_single(self, doc_id: str, fields: VespaDocumentFields) -> int:
+        return self.index.update_single(doc_id, fields)

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

@@ -1,16 +1,19 @@
-from datetime import datetime
+from http import HTTPStatus
 
+import httpx
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
-from pydantic import BaseModel
+from tenacity import RetryError
 
 from danswer.access.access import get_access_for_document
-from danswer.background.celery.celery_app import task_logger
+from danswer.background.celery.apps.app_base import task_logger
+from danswer.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex
 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
@@ -19,20 +22,20 @@ 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
 
-class RedisConnectorIndexingFenceData(BaseModel):
-    index_attempt_id: int
-    started: datetime | None
-    submitted: datetime
-    celery_task_id: str
+
+# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT
+LIGHT_SOFT_TIME_LIMIT = 105
+LIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15
 
 
 @shared_task(
     name="document_by_cc_pair_cleanup_task",
+    soft_time_limit=LIGHT_SOFT_TIME_LIMIT,
+    time_limit=LIGHT_TIME_LIMIT,
+    max_retries=DOCUMENT_BY_CC_PAIR_CLEANUP_MAX_RETRIES,
     bind=True,
-    soft_time_limit=45,
-    time_limit=60,
-    max_retries=3,
 )
 def document_by_cc_pair_cleanup_task(
     self: Task,
@@ -56,7 +59,7 @@ 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"document_id={document_id}")
+    task_logger.info(f"tenant={tenant_id} doc={document_id}")
 
     try:
         with get_session_with_tenant(tenant_id) as db_session:
@@ -64,17 +67,19 @@ def document_by_cc_pair_cleanup_task(
             chunks_affected = 0
 
             curr_ind_name, sec_ind_name = get_both_index_names(db_session)
-            document_index = get_default_document_index(
+            doc_index = get_default_document_index(
                 primary_index_name=curr_ind_name, secondary_index_name=sec_ind_name
             )
 
+            retry_index = RetryDocumentIndex(doc_index)
+
             count = get_document_connector_count(db_session, document_id)
             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)
+                chunks_affected = retry_index.delete_single(document_id)
                 delete_documents_complete__no_commit(
                     db_session=db_session,
                     document_ids=[document_id],
@@ -104,9 +109,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
-                )
+                chunks_affected = retry_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(
@@ -122,23 +125,58 @@ def document_by_cc_pair_cleanup_task(
             else:
                 pass
 
+            db_session.commit()
+
             task_logger.info(
-                f"tenant_id={tenant_id} "
-                f"document_id={document_id} "
+                f"tenant={tenant_id} "
+                f"doc={document_id} "
                 f"action={action} "
                 f"refcount={count} "
                 f"chunks={chunks_affected}"
             )
-            db_session.commit()
     except SoftTimeLimitExceeded:
         task_logger.info(
-            f"SoftTimeLimitExceeded exception. tenant_id={tenant_id} doc_id={document_id}"
+            f"SoftTimeLimitExceeded exception. tenant={tenant_id} doc={document_id}"
         )
-    except Exception as e:
-        task_logger.exception("Unexpected exception")
+        return False
+    except Exception as ex:
+        if isinstance(ex, RetryError):
+            task_logger.info(f"Retry failed: {ex.last_attempt.attempt_number}")
 
-        # 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)
+            # only set the inner exception if it is of type Exception
+            e_temp = ex.last_attempt.exception()
+            if isinstance(e_temp, Exception):
+                e = e_temp
+        else:
+            e = ex
+
+        if isinstance(e, httpx.HTTPStatusError):
+            if e.response.status_code == HTTPStatus.BAD_REQUEST:
+                task_logger.exception(
+                    f"Non-retryable HTTPStatusError: "
+                    f"tenant={tenant_id} "
+                    f"doc={document_id} "
+                    f"status={e.response.status_code}"
+                )
+            return False
+
+        task_logger.exception(
+            f"Unexpected exception: tenant={tenant_id} doc={document_id}"
+        )
+
+        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={tenant_id} doc={document_id}"
+            )
+            with get_session_with_tenant(tenant_id):
+                mark_document_as_modified(document_id, db_session)
+        return False
 
     return True

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

@@ -4,7 +4,9 @@ from datetime import timezone
 from http import HTTPStatus
 from typing import cast
 
+import httpx
 import redis
+from celery import Celery
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
@@ -12,18 +14,14 @@ from celery.result import AsyncResult
 from celery.states import READY_STATES
 from redis import Redis
 from sqlalchemy.orm import Session
+from tenacity import RetryError
 
 from danswer.access.access import get_access_for_document
-from danswer.background.celery.celery_app import celery_app
-from danswer.background.celery.celery_app import task_logger
+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_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.tasks import RedisConnectorIndexingFenceData
+from danswer.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex
+from danswer.background.celery.tasks.shared.tasks import LIGHT_SOFT_TIME_LIMIT
+from danswer.background.celery.tasks.shared.tasks import LIGHT_TIME_LIMIT
 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
@@ -54,11 +52,17 @@ 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.redis.redis_connector import RedisConnector
+from danswer.redis.redis_connector_credential_pair import RedisConnectorCredentialPair
+from danswer.redis.redis_connector_delete import RedisConnectorDelete
+from danswer.redis.redis_connector_index import RedisConnectorIndex
+from danswer.redis.redis_connector_prune import RedisConnectorPrune
+from danswer.redis.redis_document_set import RedisDocumentSet
 from danswer.redis.redis_pool import get_redis_client
+from danswer.redis.redis_usergroup import RedisUserGroup
 from danswer.utils.logger import setup_logger
 from danswer.utils.variable_functionality import fetch_versioned_implementation
 from danswer.utils.variable_functionality import (
@@ -76,8 +80,9 @@ logger = setup_logger()
     name="check_for_vespa_sync_task",
     soft_time_limit=JOB_TIMEOUT,
     trail=False,
+    bind=True,
 )
-def check_for_vespa_sync_task(*, tenant_id: str | None) -> None:
+def check_for_vespa_sync_task(self: Task, *, tenant_id: str | None) -> None:
     """Runs periodically to check if any document needs syncing.
     Generates sets of tasks for Celery if syncing is needed."""
 
@@ -94,49 +99,71 @@ def check_for_vespa_sync_task(*, tenant_id: str | None) -> None:
             return
 
         with get_session_with_tenant(tenant_id) as db_session:
-            try_generate_stale_document_sync_tasks(db_session, r, lock_beat, tenant_id)
+            try_generate_stale_document_sync_tasks(
+                self.app, db_session, r, lock_beat, tenant_id
+            )
 
+        # 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(
-                    document_set, db_session, r, lock_beat, tenant_id
+                    self.app, document_set_id, db_session, r, lock_beat, tenant_id
                 )
+        # endregion
 
-            # check if any user groups are not synced
-            if global_version.is_ee_version():
-                try:
-                    fetch_user_groups = fetch_versioned_implementation(
-                        "danswer.db.user_group", "fetch_user_groups"
-                    )
-
+        # check if any user groups are not synced
+        if global_version.is_ee_version():
+            try:
+                fetch_user_groups = fetch_versioned_implementation(
+                    "danswer.db.user_group", "fetch_user_groups"
+                )
+            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(
-                            usergroup, db_session, r, lock_beat, tenant_id
+                            self.app, usergroup_id, db_session, r, lock_beat, tenant_id
                         )
-                except ModuleNotFoundError:
-                    # Always exceptions on the MIT version, which is expected
-                    # We shouldn't actually get here if the ee version check works
-                    pass
 
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
         )
     except Exception:
-        task_logger.exception("Unexpected exception")
+        task_logger.exception(f"Unexpected exception: tenant={tenant_id}")
     finally:
         if lock_beat.owned():
             lock_beat.release()
 
 
 def try_generate_stale_document_sync_tasks(
-    db_session: Session, r: Redis, lock_beat: redis.lock.Lock, tenant_id: str | None
+    celery_app: Celery,
+    db_session: Session,
+    r: Redis,
+    lock_beat: redis.lock.Lock,
+    tenant_id: str | None,
 ) -> int | None:
     # the fence is up, do nothing
     if r.exists(RedisConnectorCredentialPair.get_fence_key()):
@@ -160,7 +187,7 @@ def try_generate_stale_document_sync_tasks(
     total_tasks_generated = 0
     cc_pairs = get_connector_credential_pairs(db_session)
     for cc_pair in cc_pairs:
-        rc = RedisConnectorCredentialPair(cc_pair.id)
+        rc = RedisConnectorCredentialPair(tenant_id, cc_pair.id)
         tasks_generated = rc.generate_tasks(
             celery_app, db_session, r, lock_beat, tenant_id
         )
@@ -187,7 +214,8 @@ def try_generate_stale_document_sync_tasks(
 
 
 def try_generate_document_set_sync_tasks(
-    document_set: DocumentSet,
+    celery_app: Celery,
+    document_set_id: int,
     db_session: Session,
     r: Redis,
     lock_beat: redis.lock.Lock,
@@ -195,15 +223,18 @@ def try_generate_document_set_sync_tasks(
 ) -> int | None:
     lock_beat.reacquire()
 
-    rds = RedisDocumentSet(document_set.id)
+    rds = RedisDocumentSet(tenant_id, document_set_id)
 
     # don't generate document set sync tasks if tasks are still pending
-    if r.exists(rds.fence_key):
+    if rds.fenced:
         return None
 
     # don't generate sync tasks if we're up to date
     # race condition with the monitor/cleanup function if we use a cached result!
-    db_session.refresh(document_set)
+    document_set = get_document_set_by_id(db_session, document_set_id)
+    if not document_set:
+        return None
+
     if document_set.is_up_to_date:
         return None
 
@@ -233,12 +264,13 @@ def try_generate_document_set_sync_tasks(
     )
 
     # set this only after all tasks have been added
-    r.set(rds.fence_key, tasks_generated)
+    rds.set_fence(tasks_generated)
     return tasks_generated
 
 
 def try_generate_user_group_sync_tasks(
-    usergroup: UserGroup,
+    celery_app: Celery,
+    usergroup_id: int,
     db_session: Session,
     r: Redis,
     lock_beat: redis.lock.Lock,
@@ -246,14 +278,20 @@ def try_generate_user_group_sync_tasks(
 ) -> int | None:
     lock_beat.reacquire()
 
-    rug = RedisUserGroup(usergroup.id)
-
-    # don't generate sync tasks if tasks are still pending
-    if r.exists(rug.fence_key):
+    rug = RedisUserGroup(tenant_id, usergroup_id)
+    if rug.fenced:
+        # don't generate sync tasks if tasks are still pending
         return None
 
     # race condition with the monitor/cleanup function if we use a cached result!
-    db_session.refresh(usergroup)
+    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
+
     if usergroup.is_up_to_date:
         return None
 
@@ -282,7 +320,7 @@ def try_generate_user_group_sync_tasks(
     )
 
     # set this only after all tasks have been added
-    r.set(rug.fence_key, tasks_generated)
+    rug.set_fence(tasks_generated)
     return tasks_generated
 
 
@@ -308,7 +346,7 @@ def monitor_connector_taskset(r: Redis) -> None:
 
 
 def monitor_document_set_taskset(
-    key_bytes: bytes, r: Redis, db_session: Session
+    tenant_id: str | None, 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)
@@ -318,21 +356,17 @@ def monitor_document_set_taskset(
 
     document_set_id = int(document_set_id_str)
 
-    rds = RedisDocumentSet(document_set_id)
-
-    fence_value = r.get(rds.fence_key)
-    if fence_value is None:
+    rds = RedisDocumentSet(tenant_id, document_set_id)
+    if not rds.fenced:
         return
 
-    try:
-        initial_count = int(cast(int, fence_value))
-    except ValueError:
-        task_logger.error("The value is not an integer.")
+    initial_count = rds.payload
+    if initial_count is None:
         return
 
     count = cast(int, r.scard(rds.taskset_key))
     task_logger.info(
-        f"Document set sync progress: document_set_id={document_set_id} "
+        f"Document set sync progress: document_set={document_set_id} "
         f"remaining={count} initial={initial_count}"
     )
     if count > 0:
@@ -347,46 +381,46 @@ 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 with ID: '{document_set_id}'!"
+                f"Successfully deleted document set: document_set={document_set_id}"
             )
         else:
             mark_document_set_as_synced(document_set_id, db_session)
             task_logger.info(
-                f"Successfully synced document set with ID: '{document_set_id}'!"
+                f"Successfully synced document set: document_set={document_set_id}"
             )
 
-    r.delete(rds.taskset_key)
-    r.delete(rds.fence_key)
+    rds.reset()
 
 
 def monitor_connector_deletion_taskset(
-    key_bytes: bytes, r: Redis, tenant_id: str | None
+    tenant_id: str | None, 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)
+    cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)
     if cc_pair_id_str 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)
+    redis_connector = RedisConnector(tenant_id, cc_pair_id)
 
-    fence_value = r.get(rcd.fence_key)
-    if fence_value is None:
+    fence_data = redis_connector.delete.payload
+    if not fence_data:
+        task_logger.warning(
+            f"Connector deletion - fence payload invalid: cc_pair={cc_pair_id}"
+        )
         return
 
-    try:
-        initial_count = int(cast(int, fence_value))
-    except ValueError:
-        task_logger.error("The value is not an integer.")
+    if fence_data.num_tasks is None:
+        # the fence is setting up but isn't ready yet
         return
 
-    count = cast(int, r.scard(rcd.taskset_key))
+    remaining = redis_connector.delete.get_remaining()
     task_logger.info(
-        f"Connector deletion progress: cc_pair={cc_pair_id} remaining={count} initial={initial_count}"
+        f"Connector deletion progress: cc_pair={cc_pair_id} remaining={remaining} initial={fence_data.num_tasks}"
     )
-    if count > 0:
+    if remaining > 0:
         return
 
     with get_session_with_tenant(tenant_id) as db_session:
@@ -447,7 +481,7 @@ def monitor_connector_deletion_taskset(
             )
             if not connector or not len(connector.credentials):
                 task_logger.info(
-                    "Found no credentials left for connector, deleting connector"
+                    "Connector deletion - Found no credentials left for connector, deleting connector"
                 )
                 db_session.delete(connector)
             db_session.commit()
@@ -457,28 +491,28 @@ def monitor_connector_deletion_taskset(
             error_message = f"Error: {str(e)}\n\nStack Trace:\n{stack_trace}"
             add_deletion_failure_message(db_session, cc_pair_id, error_message)
             task_logger.exception(
-                f"Failed to run connector_deletion. "
+                f"Connector deletion exceptioned: "
                 f"cc_pair={cc_pair_id} connector={cc_pair.connector_id} credential={cc_pair.credential_id}"
             )
             raise e
 
     task_logger.info(
-        f"Successfully deleted cc_pair: "
+        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={initial_count}"
+        f"docs_deleted={fence_data.num_tasks}"
     )
 
-    r.delete(rcd.taskset_key)
-    r.delete(rcd.fence_key)
+    redis_connector.delete.taskset_clear()
+    redis_connector.delete.set_fence(None)
 
 
 def monitor_ccpair_pruning_taskset(
-    key_bytes: bytes, r: Redis, db_session: Session
+    tenant_id: str | None, 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)
+    cc_pair_id_str = RedisConnector.get_id_from_fence_key(fence_key)
     if cc_pair_id_str is None:
         task_logger.warning(
             f"monitor_ccpair_pruning_taskset: could not parse cc_pair_id from {fence_key}"
@@ -487,46 +521,37 @@ def monitor_ccpair_pruning_taskset(
 
     cc_pair_id = int(cc_pair_id_str)
 
-    rcp = RedisConnectorPruning(cc_pair_id)
-
-    fence_value = r.get(rcp.fence_key)
-    if fence_value is None:
+    redis_connector = RedisConnector(tenant_id, cc_pair_id)
+    if not redis_connector.prune.fenced:
         return
 
-    generator_value = r.get(rcp.generator_complete_key)
-    if generator_value is None:
+    initial = redis_connector.prune.generator_complete
+    if initial is None:
         return
 
-    try:
-        initial_count = int(cast(int, generator_value))
-    except ValueError:
-        task_logger.error("The value is not an integer.")
-        return
-
-    count = cast(int, r.scard(rcp.taskset_key))
+    remaining = redis_connector.prune.get_remaining()
     task_logger.info(
-        f"Connector pruning progress: cc_pair_id={cc_pair_id} remaining={count} initial={initial_count}"
+        f"Connector pruning progress: cc_pair={cc_pair_id} remaining={remaining} initial={initial}"
     )
-    if count > 0:
+    if remaining > 0:
         return
 
     mark_ccpair_as_pruned(int(cc_pair_id), db_session)
     task_logger.info(
-        f"Successfully pruned connector credential pair. cc_pair_id={cc_pair_id}"
+        f"Successfully pruned connector credential pair. cc_pair={cc_pair_id}"
     )
 
-    r.delete(rcp.taskset_key)
-    r.delete(rcp.generator_progress_key)
-    r.delete(rcp.generator_complete_key)
-    r.delete(rcp.fence_key)
+    redis_connector.prune.taskset_clear()
+    redis_connector.prune.generator_clear()
+    redis_connector.prune.set_fence(False)
 
 
 def monitor_ccpair_indexing_taskset(
-    key_bytes: bytes, r: Redis, db_session: Session
+    tenant_id: str | None, 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)
+    composite_id = RedisConnector.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}"
@@ -541,48 +566,37 @@ def monitor_ccpair_indexing_taskset(
     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:
+    redis_connector = RedisConnector(tenant_id, cc_pair_id)
+    redis_connector_index = redis_connector.new_index(search_settings_id)
+    if not redis_connector_index.fenced:
         return
 
-    try:
-        fence_json = fence_value.decode("utf-8")
-        fence_data = RedisConnectorIndexingFenceData.model_validate_json(
-            cast(str, fence_json)
+    payload = redis_connector_index.payload
+    if not payload:
+        return
+
+    elapsed_submitted = datetime.now(timezone.utc) - payload.submitted
+
+    progress = redis_connector_index.get_progress()
+    if progress is not None:
+        task_logger.info(
+            f"Connector indexing progress: cc_pair_id={cc_pair_id} "
+            f"search_settings_id={search_settings_id} "
+            f"progress={progress} "
+            f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
         )
-    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 payload.index_attempt_id is None or payload.celery_task_id is None:
+        # the task is still setting up
+        return
 
     # Read result state BEFORE generator_complete_key to avoid a race condition
-    result: AsyncResult = AsyncResult(fence_data.celery_task_id)
+    # never use any blocking methods on the result from inside a task!
+    result: AsyncResult = AsyncResult(payload.celery_task_id)
     result_state = result.state
 
-    generator_complete_value = r.get(rci.generator_complete_key)
-    if generator_complete_value is None:
+    status_int = redis_connector_index.get_completion()
+    if status_int 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
@@ -593,30 +607,18 @@ def monitor_ccpair_indexing_taskset(
                 f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
             )
 
-            index_attempt = get_index_attempt(db_session, fence_data.index_attempt_id)
+            index_attempt = get_index_attempt(db_session, payload.index_attempt_id)
             if index_attempt:
                 mark_attempt_failed(
-                    index_attempt=index_attempt,
+                    index_attempt_id=payload.index_attempt_id,
                     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)
+            redis_connector_index.reset()
         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."
-        )
+    status_enum = HTTPStatus(status_int)
 
     task_logger.info(
         f"Connector indexing finished: cc_pair_id={cc_pair_id} "
@@ -625,11 +627,7 @@ def monitor_ccpair_indexing_taskset(
         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)
+    redis_connector_index.reset()
 
 
 @shared_task(name="monitor_vespa_sync", soft_time_limit=300, bind=True)
@@ -641,7 +639,7 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     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
+    Returns True if the task actually did work, False if it exited early to prevent overlap
     """
     r = get_redis_client(tenant_id=tenant_id)
 
@@ -679,36 +677,9 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
             f"pruning={n_pruning}"
         )
 
-        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 + "*"):
-            monitor_connector_deletion_taskset(key_bytes, r, tenant_id)
-
+        # do some cleanup before clearing fences
+        # check the db for any outstanding index attempts
         with get_session_with_tenant(tenant_id) as db_session:
-            lock_beat.reacquire()
-            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 + "*"):
-                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 + "*"):
-                monitor_ccpair_pruning_taskset(key_bytes, r, db_session)
-
-            # do some cleanup before clearing fences
-            # check the db for any outstanding index attempts
             attempts: list[IndexAttempt] = []
             attempts.extend(
                 get_all_index_attempts_by_status(IndexingStatus.NOT_STARTED, db_session)
@@ -719,16 +690,56 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
 
             for a in attempts:
                 # if attempts exist in the db but we don't detect them in redis, mark them as failed
-                rci = RedisConnectorIndexing(
+                fence_key = RedisConnectorIndex.fence_key_with_ids(
                     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)
+                if not r.exists(fence_key):
+                    failure_reason = (
+                        f"Unknown index attempt. Might be left over from a process restart: "
+                        f"index_attempt={a.id} "
+                        f"cc_pair={a.connector_credential_pair_id} "
+                        f"search_settings={a.search_settings_id}"
+                    )
+                    task_logger.warning(failure_reason)
+                    mark_attempt_failed(a.id, 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(RedisConnectorDelete.FENCE_PREFIX + "*"):
             lock_beat.reacquire()
-            for key_bytes in r.scan_iter(RedisConnectorIndexing.FENCE_PREFIX + "*"):
-                monitor_ccpair_indexing_taskset(key_bytes, r, db_session)
+            monitor_connector_deletion_taskset(tenant_id, 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:
+                monitor_document_set_taskset(tenant_id, 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:
+                monitor_usergroup_taskset(tenant_id, key_bytes, r, db_session)
+
+        lock_beat.reacquire()
+        for key_bytes in r.scan_iter(RedisConnectorPrune.FENCE_PREFIX + "*"):
+            lock_beat.reacquire()
+            with get_session_with_tenant(tenant_id) as db_session:
+                monitor_ccpair_pruning_taskset(tenant_id, key_bytes, r, db_session)
+
+        lock_beat.reacquire()
+        for key_bytes in r.scan_iter(RedisConnectorIndex.FENCE_PREFIX + "*"):
+            lock_beat.reacquire()
+            with get_session_with_tenant(tenant_id) as db_session:
+                monitor_ccpair_indexing_taskset(tenant_id, key_bytes, r, db_session)
 
         # uncomment for debugging if needed
         # r_celery = celery_app.broker_connection().channel().client
@@ -748,22 +759,22 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
 @shared_task(
     name="vespa_metadata_sync_task",
     bind=True,
-    soft_time_limit=45,
-    time_limit=60,
+    soft_time_limit=LIGHT_SOFT_TIME_LIMIT,
+    time_limit=LIGHT_TIME_LIMIT,
     max_retries=3,
 )
 def vespa_metadata_sync_task(
     self: Task, document_id: str, tenant_id: str | None
 ) -> bool:
-    task_logger.info(f"document_id={document_id}")
-
     try:
         with get_session_with_tenant(tenant_id) as db_session:
             curr_ind_name, sec_ind_name = get_both_index_names(db_session)
-            document_index = get_default_document_index(
+            doc_index = get_default_document_index(
                 primary_index_name=curr_ind_name, secondary_index_name=sec_ind_name
             )
 
+            retry_index = RetryDocumentIndex(doc_index)
+
             doc = get_document(document_id, db_session)
             if not doc:
                 return False
@@ -785,19 +796,43 @@ def vespa_metadata_sync_task(
             )
 
             # update Vespa. OK if doc doesn't exist. Raises exception otherwise.
-            chunks_affected = document_index.update_single(document_id, fields=fields)
+            chunks_affected = retry_index.update_single(document_id, fields)
 
             # 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}"
+                f"tenant={tenant_id} doc={document_id} action=sync chunks={chunks_affected}"
             )
     except SoftTimeLimitExceeded:
-        task_logger.info(f"SoftTimeLimitExceeded exception. doc_id={document_id}")
-    except Exception as e:
-        task_logger.exception("Unexpected exception")
+        task_logger.info(
+            f"SoftTimeLimitExceeded exception. tenant={tenant_id} doc={document_id}"
+        )
+    except Exception as ex:
+        if isinstance(ex, RetryError):
+            task_logger.warning(f"Retry failed: {ex.last_attempt.attempt_number}")
+
+            # only set the inner exception if it is of type Exception
+            e_temp = ex.last_attempt.exception()
+            if isinstance(e_temp, Exception):
+                e = e_temp
+        else:
+            e = ex
+
+        if isinstance(e, httpx.HTTPStatusError):
+            if e.response.status_code == HTTPStatus.BAD_REQUEST:
+                task_logger.exception(
+                    f"Non-retryable HTTPStatusError: "
+                    f"tenant={tenant_id} "
+                    f"doc={document_id} "
+                    f"status={e.response.status_code}"
+                )
+            return False
+
+        task_logger.exception(
+            f"Unexpected exception: tenant={tenant_id} doc={document_id}"
+        )
 
         # Exponential backoff from 2^4 to 2^6 ... i.e. 16, 32, 64
         countdown = 2 ** (self.request.retries + 4)

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

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

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

@@ -0,0 +1,17 @@
+"""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

@@ -0,0 +1,17 @@
+"""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

@@ -0,0 +1,17 @@
+"""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/celery/versioned_apps/primary.py

@@ -1,9 +1,8 @@
-"""Entry point for running celery worker / celery beat."""
+"""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()
-celery_app = fetch_versioned_implementation(
-    "danswer.background.celery.celery_app", "celery_app"
+app = fetch_versioned_implementation(
+    "danswer.background.celery.apps.primary", "celery_app"
 )

backend/danswer/background/indexing/job_client.py

@@ -11,7 +11,8 @@ from typing import Any
 from typing import Literal
 from typing import Optional
 
-from danswer.db.engine import get_sqlalchemy_engine
+from danswer.configs.constants import POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME
+from danswer.db.engine import SqlEngine
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -37,7 +38,9 @@ def _initializer(
     if kwargs is None:
         kwargs = {}
 
-    get_sqlalchemy_engine().dispose(close=False)
+    logger.info("Initializing spawned worker child process.")
+    SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME)
+    SqlEngine.init_engine(pool_size=4, max_overflow=12, pool_recycle=60)
     return func(*args, **kwargs)
 
 

backend/danswer/background/indexing/run_indexing.py

@@ -1,6 +1,7 @@
 import time
 import traceback
-from collections.abc import Callable
+from abc import ABC
+from abc import abstractmethod
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
@@ -41,6 +42,19 @@ 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,
@@ -92,7 +106,7 @@ def _run_indexing(
     db_session: Session,
     index_attempt: IndexAttempt,
     tenant_id: str | None,
-    progress_callback: Callable[[int], None] | None = None,
+    callback: RunIndexingCallbackInterface | None = None,
 ) -> None:
     """
     1. Get documents which are either new or updated from specified application
@@ -104,7 +118,13 @@ def _run_indexing(
     """
     start_time = time.time()
 
+    if index_attempt.search_settings is None:
+        raise ValueError(
+            "Search settings must be set for indexing. This should not be possible."
+        )
+
     search_settings = index_attempt.search_settings
+
     index_name = search_settings.index_name
 
     # Only update cc-pair status for primary index jobs
@@ -206,6 +226,11 @@ 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)
                 if (
                     (
@@ -263,8 +288,8 @@ def _run_indexing(
                 # be inaccurate
                 db_session.commit()
 
-                if progress_callback:
-                    progress_callback(len(doc_batch))
+                if callback:
+                    callback.progress(len(doc_batch))
 
                 # This new value is updated every batch, so UI can refresh per batch update
                 update_docs_indexed(
@@ -312,7 +337,7 @@ def _run_indexing(
                 or index_attempt.status != IndexingStatus.IN_PROGRESS
             ):
                 mark_attempt_failed(
-                    index_attempt,
+                    index_attempt.id,
                     db_session,
                     failure_reason=str(e),
                     full_exception_trace=traceback.format_exc(),
@@ -347,7 +372,7 @@ def _run_indexing(
         and index_attempt_md.num_exceptions >= batch_num
     ):
         mark_attempt_failed(
-            index_attempt,
+            index_attempt.id,
             db_session,
             failure_reason="All batches exceptioned.",
         )
@@ -394,7 +419,7 @@ def run_indexing_entrypoint(
     tenant_id: str | None,
     connector_credential_pair_id: int,
     is_ee: bool = False,
-    progress_callback: Callable[[int], None] | None = None,
+    callback: RunIndexingCallbackInterface | None = None,
 ) -> None:
     try:
         if is_ee:
@@ -417,7 +442,7 @@ def run_indexing_entrypoint(
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )
 
-            _run_indexing(db_session, attempt, tenant_id, progress_callback)
+            _run_indexing(db_session, attempt, tenant_id, callback)
 
             logger.info(
                 f"Indexing finished for tenant {tenant_id}: "

backend/danswer/background/task_name_builders.py

@@ -0,0 +1,4 @@
+def name_sync_external_doc_permissions_task(
+    cc_pair_id: int, tenant_id: str | None = None
+) -> str:
+    return f"sync_external_doc_permissions_task__{cc_pair_id}"

backend/danswer/background/update.py

@@ -1,494 +0,0 @@
-# TODO(rkuo): delete after background indexing via celery is fully vetted
-# 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 import text
-# from sqlalchemy.exc import ProgrammingError
-# 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 MULTI_TENANT
-# 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.configs.constants import TENANT_ID_PREFIX
-# 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_session_with_tenant
-# 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.document_index.vespa.index import VespaIndex
-# 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 INDEXING_MODEL_SERVER_PORT
-# from shared_configs.configs import LOG_LEVEL
-# 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], tenant_id: str | None
-# ) -> 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 get_session_with_tenant(tenant_id) 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],
-#     tenant_id: str | None,
-#     timeout_hours: int = CLEANUP_INDEXING_JOBS_TIMEOUT,
-# ) -> dict[int, Future | SimpleJob]:
-#     existing_jobs_copy = existing_jobs.copy()
-#     # clean up completed jobs
-#     with get_session_with_tenant(tenant_id) 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
-#         try:
-#             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,
-#                         )
-#         except ProgrammingError:
-#             logger.debug(f"No Connector Table exists for: {tenant_id}")
-#     return existing_jobs_copy
-# def kickoff_indexing_jobs(
-#     existing_jobs: dict[int, Future | SimpleJob],
-#     client: Client | SimpleJobClient,
-#     secondary_client: Client | SimpleJobClient,
-#     tenant_id: str | None,
-# ) -> dict[int, Future | SimpleJob]:
-#     existing_jobs_copy = existing_jobs.copy()
-#     current_session = get_session_with_tenant(tenant_id)
-#     # 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 current_session 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 current_session 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 current_session 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,
-#                     tenant_id,
-#                     attempt.connector_credential_pair_id,
-#                     global_version.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,
-#                     tenant_id,
-#                     attempt.connector_credential_pair_id,
-#                     global_version.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 get_all_tenant_ids() -> list[str] | list[None]:
-#     if not MULTI_TENANT:
-#         return [None]
-#     with get_session_with_tenant(tenant_id="public") as session:
-#         result = session.execute(
-#             text(
-#                 """
-#             SELECT schema_name
-#             FROM information_schema.schemata
-#             WHERE schema_name NOT IN ('pg_catalog', 'information_schema', 'public')"""
-#             )
-#         )
-#         tenant_ids = [row[0] for row in result]
-#     valid_tenants = [
-#         tenant
-#         for tenant in tenant_ids
-#         if tenant is None or tenant.startswith(TENANT_ID_PREFIX)
-#     ]
-#     return valid_tenants
-# def update_loop(
-#     delay: int = 10,
-#     num_workers: int = NUM_INDEXING_WORKERS,
-#     num_secondary_workers: int = NUM_SECONDARY_INDEXING_WORKERS,
-# ) -> None:
-#     if not MULTI_TENANT:
-#         # We can use this function as we are certain only the public schema exists
-#         # (explicitly for the non-`MULTI_TENANT` case)
-#         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=INDEXING_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,
-#             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[str | None, 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:
-#             logger.debug(
-#                 "Found existing indexing jobs: "
-#                 f"{[(tenant_id, list(jobs.keys())) for tenant_id, jobs in existing_jobs.items()]}"
-#             )
-#         try:
-#             tenants = get_all_tenant_ids()
-#             for tenant_id in tenants:
-#                 try:
-#                     logger.debug(
-#                         f"Processing {'index attempts' if tenant_id is None else f'tenant {tenant_id}'}"
-#                     )
-#                     with get_session_with_tenant(tenant_id) as db_session:
-#                         index_to_expire = check_index_swap(db_session=db_session)
-#                         if index_to_expire and tenant_id and MULTI_TENANT:
-#                             VespaIndex.delete_entries_by_tenant_id(
-#                                 tenant_id=tenant_id,
-#                                 index_name=index_to_expire.index_name,
-#                             )
-#                         if not MULTI_TENANT:
-#                             search_settings = get_current_search_settings(db_session)
-#                             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=INDEXING_MODEL_SERVER_PORT,
-#                                 )
-#                                 warm_up_bi_encoder(embedding_model=embedding_model)
-#                                 logger.notice("First inference complete.")
-#                     tenant_jobs = existing_jobs.get(tenant_id, {})
-#                     tenant_jobs = cleanup_indexing_jobs(
-#                         existing_jobs=tenant_jobs, tenant_id=tenant_id
-#                     )
-#                     create_indexing_jobs(existing_jobs=tenant_jobs, tenant_id=tenant_id)
-#                     tenant_jobs = kickoff_indexing_jobs(
-#                         existing_jobs=tenant_jobs,
-#                         client=client_primary,
-#                         secondary_client=client_secondary,
-#                         tenant_id=tenant_id,
-#                     )
-#                     existing_jobs[tenant_id] = tenant_jobs
-#                 except Exception as e:
-#                     logger.exception(
-#                         f"Failed to process tenant {tenant_id or 'default'}: {e}"
-#                     )
-#         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/models.py

@@ -10,7 +10,7 @@ from danswer.search.enums import QueryFlow
 from danswer.search.enums import SearchType
 from danswer.search.models import RetrievalDocs
 from danswer.search.models import SearchResponse
-from danswer.tools.custom.base_tool_types import ToolResultType
+from danswer.tools.tool_implementations.custom.base_tool_types import ToolResultType
 
 
 class LlmDoc(BaseModel):
@@ -156,7 +156,7 @@ class QAResponse(SearchResponse, DanswerAnswer):
     error_msg: str | None = None
 
 
-class ImageGenerationDisplay(BaseModel):
+class FileChatDisplay(BaseModel):
     file_ids: list[str]
 
 
@@ -170,7 +170,7 @@ AnswerQuestionPossibleReturn = (
     | DanswerQuotes
     | CitationInfo
     | DanswerContexts
-    | ImageGenerationDisplay
+    | FileChatDisplay
     | CustomToolResponse
     | StreamingError
     | StreamStopInfo

backend/danswer/chat/personas.yaml

@@ -41,6 +41,15 @@ personas:
     icon_color: "#6FB1FF"
     display_priority: 1
     is_visible: true
+    starter_messages:
+      - name: "Give me an overview of what's here"
+        message: "Sample some documents and tell me what you find."
+      - name: "Use AI to solve a work related problem"
+        message: "Ask me what problem I would like to solve, then search the knowledge base to help me find a solution."
+      - name: "Find updates on a topic of interest"
+        message: "Once I provide a topic, retrieve related documents and tell me when there was last activity on the topic if available."
+      - name: "Surface contradictions"
+        message: "Have me choose a subject. Once I have provided it, check against the knowledge base and point out any inconsistencies. For all your following responses, focus on identifying contradictions."
 
   - id: 1
     name: "General"
@@ -57,6 +66,15 @@ personas:
     icon_color: "#FF6F6F"
     display_priority: 0
     is_visible: true
+    starter_messages:
+      - name: "Summarize a document"
+        message: "If I have provided a document please summarize it for me. If not, please ask me to upload a document either by dragging it into the input bar or clicking the +file icon."
+      - name: "Help me with coding"
+        message: 'Write me a "Hello World" script in 5 random languages to show off the functionality.'
+      - name: "Draft a professional email"
+        message: "Help me craft a professional email. Let's establish the context and the anticipated outcomes of the email before proposing a draft."
+      - name: "Learn something new"
+        message: "What is the difference between a Gantt chart, a Burndown chart and a Kanban board?"
 
   - id: 2
     name: "Paraphrase"
@@ -73,7 +91,15 @@ personas:
     icon_color: "#6FFF8D"
     display_priority: 2
     is_visible: false
-
+    starter_messages:
+      - name: "Document Search"
+        message: "Hi! Could you help me find information about our team structure and reporting lines from our internal documents?"
+      - name: "Process Verification"
+        message: "Hello! I need to understand our project approval process. Could you find the exact steps from our documentation?"
+      - name: "Technical Documentation"
+        message: "Hi there! I'm looking for information about our deployment procedures. Can you find the specific steps from our technical guides?"
+      - name: "Policy Reference"
+        message: "Hello! Could you help me find our official guidelines about client communication? I need the exact wording from our documentation."
 
   - id: 3
     name: "Art"
@@ -86,8 +112,17 @@ personas:
     llm_filter_extraction: false
     recency_bias: "no_decay"
     document_sets: []
-    icon_shape: 234124 
+    icon_shape: 234124
     icon_color: "#9B59B6"
-    image_generation: true 
+    image_generation: true
     display_priority: 3
     is_visible: true
+    starter_messages:
+      - name: "Create visuals for a presentation"
+        message: "Generate someone presenting a graph which clearly demonstrates an upwards trajectory."
+      - name: "Find inspiration for a marketing campaign"
+        message: "Generate an image of two happy individuals sipping on a soda drink in a glass bottle."
+      - name: "Visualize a product design"
+        message: "I want to add a search bar to my Iphone app. Generate me generic examples of how other apps implement this."
+      - name: "Generate a humorous image response"
+        message: "My teammate just made a silly mistake and I want to respond with a facepalm. Can you generate me one?"

backend/danswer/chat/process_message.py

@@ -11,23 +11,18 @@ from danswer.chat.models import AllCitations
 from danswer.chat.models import CitationInfo
 from danswer.chat.models import CustomToolResponse
 from danswer.chat.models import DanswerAnswerPiece
+from danswer.chat.models import FileChatDisplay
 from danswer.chat.models import FinalUsedContextDocsResponse
-from danswer.chat.models import ImageGenerationDisplay
 from danswer.chat.models import LLMRelevanceFilterResponse
 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.chat.models import StreamStopInfo
 from danswer.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from danswer.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
 from danswer.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT
 from danswer.configs.constants import MessageType
-from danswer.configs.model_configs import GEN_AI_TEMPERATURE
 from danswer.db.chat import attach_files_to_chat_message
 from danswer.db.chat import create_db_search_doc
 from danswer.db.chat import create_new_chat_message
@@ -40,7 +35,6 @@ from danswer.db.chat import reserve_message_id
 from danswer.db.chat import translate_db_message_to_chat_message_detail
 from danswer.db.chat import translate_db_search_doc_to_server_search_doc
 from danswer.db.engine import get_session_context_manager
-from danswer.db.llm import fetch_existing_llm_providers
 from danswer.db.models import SearchDoc as DbSearchDoc
 from danswer.db.models import ToolCall
 from danswer.db.models import User
@@ -60,14 +54,13 @@ from danswer.llm.answering.models import PromptConfig
 from danswer.llm.exceptions import GenAIDisabledException
 from danswer.llm.factory import get_llms_for_persona
 from danswer.llm.factory import get_main_llm_from_tuple
-from danswer.llm.interfaces import LLMConfig
 from danswer.llm.utils import litellm_exception_to_error_msg
 from danswer.natural_language_processing.utils import get_tokenizer
-from danswer.search.enums import LLMEvaluationType
 from danswer.search.enums import OptionalSearchSetting
 from danswer.search.enums import QueryFlow
 from danswer.search.enums import SearchType
 from danswer.search.models import InferenceSection
+from danswer.search.models import RetrievalDetails
 from danswer.search.retrieval.search_runner import inference_sections_from_ids
 from danswer.search.utils import chunks_or_sections_to_search_docs
 from danswer.search.utils import dedupe_documents
@@ -76,36 +69,48 @@ from danswer.search.utils import relevant_sections_to_indices
 from danswer.server.query_and_chat.models import ChatMessageDetail
 from danswer.server.query_and_chat.models import CreateChatMessageRequest
 from danswer.server.utils import get_json_line
-from danswer.tools.built_in_tools import get_built_in_tool_by_id
-from danswer.tools.custom.custom_tool import (
-    build_custom_tools_from_openapi_schema_and_headers,
-)
-from danswer.tools.custom.custom_tool import CUSTOM_TOOL_RESPONSE_ID
-from danswer.tools.custom.custom_tool import CustomToolCallSummary
 from danswer.tools.force import ForceUseTool
-from danswer.tools.images.image_generation_tool import IMAGE_GENERATION_RESPONSE_ID
-from danswer.tools.images.image_generation_tool import ImageGenerationResponse
-from danswer.tools.images.image_generation_tool import ImageGenerationTool
-from danswer.tools.internet_search.internet_search_tool import (
+from danswer.tools.models import ToolResponse
+from danswer.tools.tool import Tool
+from danswer.tools.tool_constructor import construct_tools
+from danswer.tools.tool_constructor import CustomToolConfig
+from danswer.tools.tool_constructor import ImageGenerationToolConfig
+from danswer.tools.tool_constructor import InternetSearchToolConfig
+from danswer.tools.tool_constructor import SearchToolConfig
+from danswer.tools.tool_implementations.custom.custom_tool import (
+    CUSTOM_TOOL_RESPONSE_ID,
+)
+from danswer.tools.tool_implementations.custom.custom_tool import CustomToolCallSummary
+from danswer.tools.tool_implementations.images.image_generation_tool import (
+    IMAGE_GENERATION_RESPONSE_ID,
+)
+from danswer.tools.tool_implementations.images.image_generation_tool import (
+    ImageGenerationResponse,
+)
+from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
     INTERNET_SEARCH_RESPONSE_ID,
 )
-from danswer.tools.internet_search.internet_search_tool import (
+from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
     internet_search_response_to_search_docs,
 )
-from danswer.tools.internet_search.internet_search_tool import InternetSearchResponse
-from danswer.tools.internet_search.internet_search_tool import InternetSearchTool
-from danswer.tools.models import DynamicSchemaInfo
-from danswer.tools.search.search_tool import FINAL_CONTEXT_DOCUMENTS_ID
-from danswer.tools.search.search_tool import SEARCH_RESPONSE_SUMMARY_ID
-from danswer.tools.search.search_tool import SearchResponseSummary
-from danswer.tools.search.search_tool import SearchTool
-from danswer.tools.search.search_tool import SECTION_RELEVANCE_LIST_ID
-from danswer.tools.tool import Tool
-from danswer.tools.tool import ToolResponse
+from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
+    InternetSearchResponse,
+)
+from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
+    InternetSearchTool,
+)
+from danswer.tools.tool_implementations.search.search_tool import (
+    FINAL_CONTEXT_DOCUMENTS_ID,
+)
+from danswer.tools.tool_implementations.search.search_tool import (
+    SEARCH_RESPONSE_SUMMARY_ID,
+)
+from danswer.tools.tool_implementations.search.search_tool import SearchResponseSummary
+from danswer.tools.tool_implementations.search.search_tool import SearchTool
+from danswer.tools.tool_implementations.search.search_tool import (
+    SECTION_RELEVANCE_LIST_ID,
+)
 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
 
@@ -256,10 +261,11 @@ ChatPacket = (
     | DanswerAnswerPiece
     | AllCitations
     | CitationInfo
-    | ImageGenerationDisplay
+    | FileChatDisplay
     | CustomToolResponse
     | MessageSpecificCitations
     | MessageResponseIDInfo
+    | StreamStopInfo
 )
 ChatPacketStream = Iterator[ChatPacket]
 
@@ -275,7 +281,6 @@ def stream_chat_message_objects(
     max_document_percentage: float = CHAT_TARGET_CHUNK_PERCENTAGE,
     # if specified, uses the last user message and does not create a new user message based
     # 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,
@@ -287,6 +292,9 @@ def stream_chat_message_objects(
     3. [always] A set of streamed LLM tokens or an error anywhere along the line if something fails
     4. [always] Details on the final AI response message that is created
     """
+    use_existing_user_message = new_msg_req.use_existing_user_message
+    existing_assistant_message_id = new_msg_req.existing_assistant_message_id
+
     # Currently surrounding context is not supported for chat
     # Chat is already token heavy and harder for the model to process plus it would roll history over much faster
     new_msg_req.chunks_above = 0
@@ -408,12 +416,20 @@ def stream_chat_message_objects(
             final_msg, history_msgs = create_chat_chain(
                 chat_session_id=chat_session_id, db_session=db_session
             )
-            if final_msg.message_type != MessageType.USER:
-                raise RuntimeError(
-                    "The last message was not a user message. Cannot call "
-                    "`stream_chat_message_objects` with `is_regenerate=True` "
-                    "when the last message is not a user message."
-                )
+            if existing_assistant_message_id is None:
+                if final_msg.message_type != MessageType.USER:
+                    raise RuntimeError(
+                        "The last message was not a user message. Cannot call "
+                        "`stream_chat_message_objects` with `is_regenerate=True` "
+                        "when the last message is not a user message."
+                    )
+            else:
+                if final_msg.id != existing_assistant_message_id:
+                    raise RuntimeError(
+                        "The last message was not the existing assistant message. "
+                        f"Final message id: {final_msg.id}, "
+                        f"existing assistant message id: {existing_assistant_message_id}"
+                    )
 
         # Disable Query Rephrasing for the first message
         # This leads to a better first response since the LLM rephrasing the question
@@ -484,13 +500,19 @@ def stream_chat_message_objects(
                 ),
                 max_window_percentage=max_document_percentage,
             )
-        reserved_message_id = reserve_message_id(
-            db_session=db_session,
-            chat_session_id=chat_session_id,
-            parent_message=user_message.id
-            if user_message is not None
-            else parent_message.id,
-            message_type=MessageType.ASSISTANT,
+
+        # we don't need to reserve a message id if we're using an existing assistant message
+        reserved_message_id = (
+            final_msg.id
+            if existing_assistant_message_id is not None
+            else reserve_message_id(
+                db_session=db_session,
+                chat_session_id=chat_session_id,
+                parent_message=user_message.id
+                if user_message is not None
+                else parent_message.id,
+                message_type=MessageType.ASSISTANT,
+            )
         )
         yield MessageResponseIDInfo(
             user_message_id=user_message.id if user_message else None,
@@ -505,7 +527,13 @@ def stream_chat_message_objects(
         partial_response = partial(
             create_new_chat_message,
             chat_session_id=chat_session_id,
-            parent_message=final_msg,
+            # if we're using an existing assistant message, then this will just be an
+            # update operation, in which case the parent should be the parent of
+            # the latest. If we're creating a new assistant message, then the parent
+            # should be the latest message (latest user message)
+            parent_message=(
+                final_msg if existing_assistant_message_id is None else parent_message
+            ),
             prompt_id=prompt_id,
             overridden_model=overridden_model,
             # message=,
@@ -517,6 +545,7 @@ def stream_chat_message_objects(
             # reference_docs=,
             db_session=db_session,
             commit=False,
+            reserved_message_id=reserved_message_id,
         )
 
         if not final_msg.prompt:
@@ -532,147 +561,53 @@ def stream_chat_message_objects(
             if not persona
             else PromptConfig.from_model(persona.prompts[0])
         )
+        answer_style_config = AnswerStyleConfig(
+            citation_config=CitationConfig(
+                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,
+        )
 
-        # find out what tools to use
-        search_tool: SearchTool | None = None
-        tool_dict: dict[int, list[Tool]] = {}  # tool_id to tool
-        for db_tool_model in persona.tools:
-            # handle in-code tools specially
-            if db_tool_model.in_code_tool_id:
-                tool_cls = get_built_in_tool_by_id(db_tool_model.id, db_session)
-                if tool_cls.__name__ == SearchTool.__name__ and not latest_query_files:
-                    search_tool = SearchTool(
-                        db_session=db_session,
-                        user=user,
-                        persona=persona,
-                        retrieval_options=retrieval_options,
-                        prompt_config=prompt_config,
-                        llm=llm,
-                        fast_llm=fast_llm,
-                        pruning_config=document_pruning_config,
-                        selected_sections=selected_sections,
-                        chunks_above=new_msg_req.chunks_above,
-                        chunks_below=new_msg_req.chunks_below,
-                        full_doc=new_msg_req.full_doc,
-                        evaluation_type=LLMEvaluationType.BASIC
-                        if persona.llm_relevance_filter
-                        else LLMEvaluationType.SKIP,
-                    )
-                    tool_dict[db_tool_model.id] = [search_tool]
-                elif tool_cls.__name__ == ImageGenerationTool.__name__:
-                    img_generation_llm_config: LLMConfig | None = None
-                    if (
-                        llm
-                        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,
-                        )
-                    else:
-                        llm_providers = fetch_existing_llm_providers(db_session)
-                        openai_provider = next(
-                            iter(
-                                [
-                                    llm_provider
-                                    for llm_provider in llm_providers
-                                    if llm_provider.provider == "openai"
-                                ]
-                            ),
-                            None,
-                        )
-                        if not openai_provider or not openai_provider.api_key:
-                            raise ValueError(
-                                "Image generation tool requires an OpenAI API key"
-                            )
-                        img_generation_llm_config = LLMConfig(
-                            model_provider=openai_provider.provider,
-                            model_name="dall-e-3",
-                            temperature=GEN_AI_TEMPERATURE,
-                            api_key=openai_provider.api_key,
-                            api_base=openai_provider.api_base,
-                            api_version=openai_provider.api_version,
-                        )
-                    tool_dict[db_tool_model.id] = [
-                        ImageGenerationTool(
-                            api_key=cast(str, img_generation_llm_config.api_key),
-                            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__:
-                    bing_api_key = BING_API_KEY
-                    if not bing_api_key:
-                        raise ValueError(
-                            "Internet search tool requires a Bing API key, please contact your Danswer admin to get it added!"
-                        )
-                    tool_dict[db_tool_model.id] = [
-                        InternetSearchTool(api_key=bing_api_key)
-                    ]
-
-                continue
-
-            # handle all custom tools
-            if db_tool_model.openapi_schema:
-                tool_dict[db_tool_model.id] = cast(
-                    list[Tool],
-                    build_custom_tools_from_openapi_schema_and_headers(
-                        db_tool_model.openapi_schema,
-                        dynamic_schema_info=DynamicSchemaInfo(
-                            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 {}
-                            )
-                        ),
-                    ),
-                )
-
+        tool_dict = construct_tools(
+            persona=persona,
+            prompt_config=prompt_config,
+            db_session=db_session,
+            user=user,
+            llm=llm,
+            fast_llm=fast_llm,
+            search_tool_config=SearchToolConfig(
+                answer_style_config=answer_style_config,
+                document_pruning_config=document_pruning_config,
+                retrieval_options=retrieval_options or RetrievalDetails(),
+                selected_sections=selected_sections,
+                chunks_above=new_msg_req.chunks_above,
+                chunks_below=new_msg_req.chunks_below,
+                full_doc=new_msg_req.full_doc,
+                latest_query_files=latest_query_files,
+            ),
+            internet_search_tool_config=InternetSearchToolConfig(
+                answer_style_config=answer_style_config,
+            ),
+            image_generation_tool_config=ImageGenerationToolConfig(
+                additional_headers=litellm_additional_headers,
+            ),
+            custom_tool_config=CustomToolConfig(
+                chat_session_id=chat_session_id,
+                message_id=user_message.id if user_message else None,
+                additional_headers=custom_tool_additional_headers,
+            ),
+        )
         tools: list[Tool] = []
         for tool_list in tool_dict.values():
             tools.extend(tool_list)
 
-        # factor in tool definition size when pruning
-        document_pruning_config.tool_num_tokens = compute_all_tool_tokens(
-            tools, llm_tokenizer
-        )
-        document_pruning_config.using_tool_message = explicit_tool_calling_supported(
-            llm_provider, llm_model_name
-        )
-
         # LLM prompt building, response capturing, etc.
         answer = Answer(
             is_connected=is_connected,
             question=final_msg.message,
             latest_query_files=latest_query_files,
-            answer_style_config=AnswerStyleConfig(
-                citation_config=CitationConfig(
-                    all_docs_useful=selected_db_search_docs is not None
-                ),
-                document_pruning_config=document_pruning_config,
-            ),
+            answer_style_config=answer_style_config,
             prompt_config=prompt_config,
             llm=(
                 llm
@@ -740,7 +675,6 @@ def stream_chat_message_objects(
                         yield LLMRelevanceFilterResponse(
                             llm_selected_doc_indices=llm_indices
                         )
-
                 elif packet.id == FINAL_CONTEXT_DOCUMENTS_ID:
                     yield FinalUsedContextDocsResponse(
                         final_context_docs=packet.response
@@ -758,7 +692,7 @@ def stream_chat_message_objects(
                         FileDescriptor(id=str(file_id), type=ChatFileType.IMAGE)
                         for file_id in file_ids
                     ]
-                    yield ImageGenerationDisplay(
+                    yield FileChatDisplay(
                         file_ids=[str(file_id) for file_id in file_ids]
                     )
                 elif packet.id == INTERNET_SEARCH_RESPONSE_ID:
@@ -772,11 +706,32 @@ def stream_chat_message_objects(
                     yield qa_docs_response
                 elif packet.id == CUSTOM_TOOL_RESPONSE_ID:
                     custom_tool_response = cast(CustomToolCallSummary, packet.response)
-                    yield CustomToolResponse(
-                        response=custom_tool_response.tool_result,
-                        tool_name=custom_tool_response.tool_name,
-                    )
 
+                    if (
+                        custom_tool_response.response_type == "image"
+                        or custom_tool_response.response_type == "csv"
+                    ):
+                        file_ids = custom_tool_response.tool_result.file_ids
+                        ai_message_files = [
+                            FileDescriptor(
+                                id=str(file_id),
+                                type=ChatFileType.IMAGE
+                                if custom_tool_response.response_type == "image"
+                                else ChatFileType.CSV,
+                            )
+                            for file_id in file_ids
+                        ]
+                        yield FileChatDisplay(
+                            file_ids=[str(file_id) for file_id in file_ids]
+                        )
+                    else:
+                        yield CustomToolResponse(
+                            response=custom_tool_response.tool_result,
+                            tool_name=custom_tool_response.tool_name,
+                        )
+
+            elif isinstance(packet, StreamStopInfo):
+                pass
             else:
                 if isinstance(packet, ToolCallFinalResult):
                     tool_result = packet
@@ -806,6 +761,7 @@ def stream_chat_message_objects(
 
     # Post-LLM answer processing
     try:
+        logger.debug("Post-LLM answer processing")
         message_specific_citations: MessageSpecificCitations | None = None
         if reference_db_search_docs:
             message_specific_citations = _translate_citations(
@@ -821,7 +777,6 @@ def stream_chat_message_objects(
                 tool_name_to_tool_id[tool.name] = tool_id
 
         gen_ai_response_message = partial_response(
-            reserved_message_id=reserved_message_id,
             message=answer.llm_answer,
             rephrased_query=(
                 qa_docs_response.rephrased_query if qa_docs_response else None
@@ -829,21 +784,21 @@ def stream_chat_message_objects(
             reference_docs=reference_db_search_docs,
             files=ai_message_files,
             token_count=len(llm_tokenizer_encode_func(answer.llm_answer)),
-            citations=message_specific_citations.citation_map
-            if message_specific_citations
-            else None,
+            citations=(
+                message_specific_citations.citation_map
+                if message_specific_citations
+                else None
+            ),
             error=None,
-            tool_calls=(
-                [
-                    ToolCall(
-                        tool_id=tool_name_to_tool_id[tool_result.tool_name],
-                        tool_name=tool_result.tool_name,
-                        tool_arguments=tool_result.tool_args,
-                        tool_result=tool_result.tool_result,
-                    )
-                ]
+            tool_call=(
+                ToolCall(
+                    tool_id=tool_name_to_tool_id[tool_result.tool_name],
+                    tool_name=tool_result.tool_name,
+                    tool_arguments=tool_result.tool_args,
+                    tool_result=tool_result.tool_result,
+                )
                 if tool_result
-                else []
+                else None
             ),
         )
 
@@ -867,7 +822,6 @@ def stream_chat_message_objects(
 def stream_chat_message(
     new_msg_req: CreateChatMessageRequest,
     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,
@@ -877,7 +831,6 @@ def stream_chat_message(
             new_msg_req=new_msg_req,
             user=user,
             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,

backend/danswer/chat/prompts.yaml

@@ -9,19 +9,19 @@ prompts:
     system: >
       You are a question answering system that is constantly learning and improving.
       The current date is DANSWER_DATETIME_REPLACEMENT.
-      
+
       You can process and comprehend vast amounts of text and utilize this knowledge to provide
       grounded, accurate, and concise answers to diverse queries.
-      
+
       You always clearly communicate ANY UNCERTAINTY in your answer.
     # Task Prompt (as shown in UI)
     task: >
       Answer my query based on the documents provided.
       The documents may not all be relevant, ignore any documents that are not directly relevant
       to the most recent user query.
-      
+
       I have not read or seen any of the documents and do not want to read them.
-      
+
       If there are no relevant documents, refer to the chat history and your internal knowledge.
     # Inject a statement at the end of system prompt to inform the LLM of the current date/time
     # If the DANSWER_DATETIME_REPLACEMENT is set, the date/time is inserted there instead
@@ -30,21 +30,21 @@ prompts:
     # Prompts the LLM to include citations in the for [1], [2] etc.
     # which get parsed to match the passed in sources
     include_citations: true
-  
+
   - name: "ImageGeneration"
-    description: "Generates images based on user prompts!"
+    description: "Generates images from user descriptions!"
     system: >
-      You are an advanced image generation system capable of creating diverse and detailed images.
-      
-      You can interpret user prompts and generate high-quality, creative images that match their descriptions.
-      
-      You always strive to create safe and appropriate content, avoiding any harmful or offensive imagery.
+      You are an AI image generation assistant. Your role is to create high-quality images based on user descriptions.
+
+      For appropriate requests, you will generate an image that matches the user's requirements.
+      For inappropriate or unsafe requests, you will politely decline and explain why the request cannot be fulfilled.
+
+      You aim to be helpful while maintaining appropriate content standards.
     task: >
-      Generate an image based on the user's description.
-      
-      Provide a detailed description of the generated image, including key elements, colors, and composition.
-      
-      If the request is not possible or appropriate, explain why and suggest alternatives.
+      Based on the user's description, create a high-quality image that accurately reflects their request. 
+      Pay close attention to the specified details, styles, and desired elements.
+
+      If the request is not appropriate or cannot be fulfilled, explain why and suggest alternatives.
     datetime_aware: true
     include_citations: false
 
@@ -64,14 +64,13 @@ prompts:
     datetime_aware: true
     include_citations: true
 
-
   - name: "Summarize"
     description: "Summarize relevant information from retrieved context!"
     system: >
       You are a text summarizing assistant that highlights the most important knowledge from the
       context provided, prioritizing the information that relates to the user query.
       The current date is DANSWER_DATETIME_REPLACEMENT.
-      
+
       You ARE NOT creative and always stick to the provided documents.
       If there are no documents, refer to the conversation history.
 
@@ -84,7 +83,6 @@ prompts:
     datetime_aware: true
     include_citations: true
 
-
   - name: "Paraphrase"
     description: "Recites information from retrieved context! Least creative but most safe!"
     system: >
@@ -92,10 +90,10 @@ prompts:
       The current date is DANSWER_DATETIME_REPLACEMENT.
 
       You only provide quotes that are EXACT substrings from provided documents!
-      
+
       If there are no documents provided,
       simply tell the user that there are no documents to reference.
-      
+
       You NEVER generate new text or phrases outside of the citation.
       DO NOT explain your responses, only provide the quotes and NOTHING ELSE.
     task: >

backend/danswer/configs/app_configs.py

@@ -43,6 +43,9 @@ 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
@@ -131,7 +134,6 @@ try:
 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"
@@ -161,6 +163,17 @@ try:
 except ValueError:
     POSTGRES_POOL_RECYCLE = POSTGRES_POOL_RECYCLE_DEFAULT
 
+# Experimental setting to control idle transactions
+POSTGRES_IDLE_SESSIONS_TIMEOUT_DEFAULT = 0  # milliseconds
+try:
+    POSTGRES_IDLE_SESSIONS_TIMEOUT = int(
+        os.environ.get(
+            "POSTGRES_IDLE_SESSIONS_TIMEOUT", POSTGRES_IDLE_SESSIONS_TIMEOUT_DEFAULT
+        )
+    )
+except ValueError:
+    POSTGRES_IDLE_SESSIONS_TIMEOUT = POSTGRES_IDLE_SESSIONS_TIMEOUT_DEFAULT
+
 REDIS_SSL = os.getenv("REDIS_SSL", "").lower() == "true"
 REDIS_HOST = os.environ.get("REDIS_HOST") or "localhost"
 REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))
@@ -198,6 +211,41 @@ 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
 #####
@@ -214,9 +262,6 @@ ENABLED_CONNECTOR_TYPES = os.environ.get("ENABLED_CONNECTOR_TYPES") or ""
 # for some connectors
 ENABLE_EXPENSIVE_EXPERT_CALLS = False
 
-GOOGLE_DRIVE_INCLUDE_SHARED = False
-GOOGLE_DRIVE_FOLLOW_SHORTCUTS = False
-GOOGLE_DRIVE_ONLY_ORG_PUBLIC = False
 
 # TODO these should be available for frontend configuration, via advanced options expandable
 WEB_CONNECTOR_IGNORED_CLASSES = os.environ.get(
@@ -253,12 +298,6 @@ 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)
@@ -406,7 +445,7 @@ 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")
+VESPA_REQUEST_TIMEOUT = int(os.environ.get("VESPA_REQUEST_TIMEOUT") or "15")
 
 SYSTEM_RECURSION_LIMIT = int(os.environ.get("SYSTEM_RECURSION_LIMIT") or "1000")
 
@@ -430,20 +469,12 @@ 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"
-
 # 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
@@ -458,3 +489,17 @@ CONTROL_PLANE_API_BASE_URL = os.environ.get(
 
 # JWT configuration
 JWT_ALGORITHM = "HS256"
+
+# Super Users
+SUPER_USERS = json.loads(os.environ.get("SUPER_USERS", '["pablo@danswer.ai"]'))
+SUPER_CLOUD_API_KEY = os.environ.get("SUPER_CLOUD_API_KEY", "api_key")
+
+
+#####
+# API Key Configs
+#####
+# refers to the rounds described here: https://passlib.readthedocs.io/en/stable/lib/passlib.hash.sha256_crypt.html
+_API_KEY_HASH_ROUNDS_RAW = os.environ.get("API_KEY_HASH_ROUNDS")
+API_KEY_HASH_ROUNDS = (
+    int(_API_KEY_HASH_ROUNDS_RAW) if _API_KEY_HASH_ROUNDS_RAW else None
+)

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"
@@ -46,7 +43,6 @@ 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__"
@@ -71,6 +67,7 @@ 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
@@ -128,6 +125,8 @@ class DocumentSource(str, Enum):
     OCI_STORAGE = "oci_storage"
     XENFORO = "xenforo"
     NOT_APPLICABLE = "not_applicable"
+    FRESHDESK = "freshdesk"
+    FIREFLIES = "fireflies"
 
 
 DocumentSourceRequiringTenantContext: list[DocumentSource] = [DocumentSource.FILE]
@@ -136,6 +135,7 @@ DocumentSourceRequiringTenantContext: list[DocumentSource] = [DocumentSource.FIL
 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):
@@ -226,6 +226,9 @@ class DanswerRedisLocks:
     PRUNING_LOCK_PREFIX = "da_lock:pruning"
     INDEXING_METADATA_PREFIX = "da_metadata:indexing"
 
+    SLACK_BOT_LOCK = "da_lock:slack_bot"
+    SLACK_BOT_HEARTBEAT_PREFIX = "da_heartbeat:slack_bot"
+
 
 class DanswerCeleryPriority(int, Enum):
     HIGHEST = 0

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 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
+  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
   documents.
 - Event Based connectors:
   - Connectors that listen to events and update documents accordingly.

backend/danswer/connectors/axero/connector.py

@@ -15,7 +15,6 @@ 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
@@ -24,6 +23,7 @@ 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/bookstack/connector.py

@@ -44,8 +44,6 @@ class BookstackConnector(LoadConnector, PollConnector):
         start: SecondsSinceUnixEpoch | None = None,
         end: SecondsSinceUnixEpoch | None = None,
     ) -> tuple[list[Document], int]:
-        doc_batch: list[Document] = []
-
         params = {
             "count": str(batch_size),
             "offset": str(start_ind),
@@ -63,8 +61,7 @@ class BookstackConnector(LoadConnector, PollConnector):
             )
 
         batch = bookstack_client.get(endpoint, params=params).get("data", [])
-        for item in batch:
-            doc_batch.append(transformer(bookstack_client, item))
+        doc_batch = [transformer(bookstack_client, item) for item in batch]
 
         return doc_batch, len(batch)
 

backend/danswer/connectors/clickup/connector.py

@@ -10,7 +10,6 @@ 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
@@ -19,6 +18,7 @@ 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"
@@ -210,6 +210,7 @@ if __name__ == "__main__":
             "clickup_team_id": os.environ["clickup_team_id"],
         }
     )
+
     latest_docs = clickup_connector.load_from_state()
 
     for doc in latest_docs:

backend/danswer/connectors/confluence/confluence_utils.py

@@ -1,32 +0,0 @@
-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

@@ -0,0 +1,226 @@
+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

@@ -1,219 +0,0 @@
-import math
-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
-
-
-# commenting out while we try using confluence's rate limiter instead
-# # https://developer.atlassian.com/cloud/confluence/rate-limiting/
-# 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 is used when the server doesn't hand back "Retry-After"
-#         # and we have to decide the retry delay ourselves
-#         max_delay = 30  # Atlassian uses max_delay = 30 in their examples
-
-#         # max_retry_after is used when we do get a "Retry-After" header
-#         max_retry_after = 300  # should we really cap the maximum retry delay?
-
-#         NEXT_RETRY_KEY = BaseConnector.REDIS_KEY_PREFIX + "confluence_next_retry"
-
-#         # for testing purposes, rate limiting is written to fall back to a simpler
-#         # rate limiting approach when redis is not available
-#         r = get_redis_client(tenant_id=tenant_id)
-
-#         for attempt in range(max_retries):
-#             try:
-#                 # if multiple connectors are waiting for the next attempt, there could be an issue
-#                 # where many connectors are "released" onto the server at the same time.
-#                 # That's not ideal ... but coming up with a mechanism for queueing
-#                 # all of these connectors is a bigger problem that we want to take on
-#                 # right now
-#                 try:
-#                     next_attempt = r.get(NEXT_RETRY_KEY)
-#                     if next_attempt is None:
-#                         next_attempt = 0
-#                     else:
-#                         next_attempt = int(cast(int, next_attempt))
-
-#                     # TODO: all connectors need to be interruptible moving forward
-#                     while time.monotonic() < next_attempt:
-#                         time.sleep(1)
-#                 except ConnectionError:
-#                     pass
-
-#                 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 > max_retry_after:
-#                             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..."
-#                         )
-#                         try:
-#                             r.set(
-#                                 NEXT_RETRY_KEY,
-#                                 math.ceil(time.monotonic() + retry_after),
-#                             )
-#                         except ConnectionError:
-#                             pass
-#                     else:
-#                         logger.warning(
-#                             "Rate limit hit. Retrying with exponential backoff..."
-#                         )
-#                         delay = min(starting_delay * (backoff**attempt), max_delay)
-#                         delay_until = math.ceil(time.monotonic() + delay)
-
-#                         try:
-#                             r.set(NEXT_RETRY_KEY, delay_until)
-#                         except ConnectionError:
-#                             while time.monotonic() < delay_until:
-#                                 time.sleep(1)
-#                 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)
-#                 delay_until = math.ceil(time.monotonic() + delay)
-#                 try:
-#                     r.set(NEXT_RETRY_KEY, delay_until)
-#                 except ConnectionError:
-#                     while time.monotonic() < delay_until:
-#                         time.sleep(1)
-
-#                 if attempt == max_retries - 1:
-#                     raise e
-
-#     return cast(F, wrapped_call)
-
-
-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 make_confluence_call_handle_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)

backend/danswer/connectors/confluence/utils.py

@@ -0,0 +1,214 @@
+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/miscellaneous_utils.py

@@ -11,6 +11,10 @@ from danswer.connectors.models import BasicExpertInfo
 from danswer.utils.text_processing import is_valid_email
 
 
+T = TypeVar("T")
+U = TypeVar("U")
+
+
 def datetime_to_utc(dt: datetime) -> datetime:
     if dt.tzinfo is None or dt.tzinfo.utcoffset(dt) is None:
         dt = dt.replace(tzinfo=timezone.utc)
@@ -19,7 +23,16 @@ def datetime_to_utc(dt: datetime) -> datetime:
 
 
 def time_str_to_utc(datetime_str: str) -> datetime:
-    dt = parse(datetime_str)
+    try:
+        dt = parse(datetime_str)
+    except ValueError:
+        # Handle malformed timezone by attempting to fix common format issues
+        if "0000" in datetime_str:
+            # Convert "0000" to "+0000" for proper timezone parsing
+            fixed_dt_str = datetime_str.replace(" 0000", " +0000")
+            dt = parse(fixed_dt_str)
+        else:
+            raise
     return datetime_to_utc(dt)
 
 
@@ -49,10 +62,6 @@ def get_experts_stores_representations(
     return [owner for owner in reps if owner is not None]
 
 
-T = TypeVar("T")
-U = TypeVar("U")
-
-
 def process_in_batches(
     objects: list[T], process_function: Callable[[T], U], batch_size: int
 ) -> Iterator[list[U]]:

backend/danswer/connectors/discourse/connector.py

@@ -14,7 +14,6 @@ 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
@@ -24,6 +23,7 @@ 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,7 +11,6 @@ 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
@@ -22,6 +21,7 @@ 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

@@ -16,6 +16,8 @@ from danswer.connectors.discourse.connector import DiscourseConnector
 from danswer.connectors.document360.connector import Document360Connector
 from danswer.connectors.dropbox.connector import DropboxConnector
 from danswer.connectors.file.connector import LocalFileConnector
+from danswer.connectors.fireflies.connector import FirefliesConnector
+from danswer.connectors.freshdesk.connector import FreshdeskConnector
 from danswer.connectors.github.connector import GithubConnector
 from danswer.connectors.gitlab.connector import GitlabConnector
 from danswer.connectors.gmail.connector import GmailConnector
@@ -34,7 +36,6 @@ from danswer.connectors.mediawiki.wiki import MediaWikiConnector
 from danswer.connectors.models import InputType
 from danswer.connectors.notion.connector import NotionConnector
 from danswer.connectors.productboard.connector import ProductboardConnector
-from danswer.connectors.requesttracker.connector import RequestTrackerConnector
 from danswer.connectors.salesforce.connector import SalesforceConnector
 from danswer.connectors.sharepoint.connector import SharepointConnector
 from danswer.connectors.slab.connector import SlabConnector
@@ -64,7 +65,7 @@ def identify_connector_class(
         DocumentSource.SLACK: {
             InputType.LOAD_STATE: SlackLoadConnector,
             InputType.POLL: SlackPollConnector,
-            InputType.PRUNE: SlackPollConnector,
+            InputType.SLIM_RETRIEVAL: SlackPollConnector,
         },
         DocumentSource.GITHUB: GithubConnector,
         DocumentSource.GMAIL: GmailConnector,
@@ -77,7 +78,6 @@ def identify_connector_class(
         DocumentSource.SLAB: SlabConnector,
         DocumentSource.NOTION: NotionConnector,
         DocumentSource.ZULIP: ZulipConnector,
-        DocumentSource.REQUESTTRACKER: RequestTrackerConnector,
         DocumentSource.GURU: GuruConnector,
         DocumentSource.LINEAR: LinearConnector,
         DocumentSource.HUBSPOT: HubSpotConnector,
@@ -101,6 +101,8 @@ def identify_connector_class(
         DocumentSource.GOOGLE_CLOUD_STORAGE: BlobStorageConnector,
         DocumentSource.OCI_STORAGE: BlobStorageConnector,
         DocumentSource.XENFORO: XenforoConnector,
+        DocumentSource.FRESHDESK: FreshdeskConnector,
+        DocumentSource.FIREFLIES: FirefliesConnector,
     }
     connector_by_source = connector_map.get(source, {})
 

backend/danswer/connectors/file/connector.py

@@ -10,7 +10,6 @@ from sqlalchemy.orm import Session
 
 from danswer.configs.app_configs import INDEX_BATCH_SIZE
 from danswer.configs.constants import DocumentSource
-from danswer.configs.constants import POSTGRES_DEFAULT_SCHEMA
 from danswer.connectors.cross_connector_utils.miscellaneous_utils import time_str_to_utc
 from danswer.connectors.interfaces import GenerateDocumentsOutput
 from danswer.connectors.interfaces import LoadConnector
@@ -28,7 +27,8 @@ 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
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 
 logger = setup_logger()
 
@@ -123,9 +123,13 @@ def _process_file(
             "filename",
             "file_display_name",
             "title",
+            "connector_type",
         ]
     }
 
+    source_type_str = all_metadata.get("connector_type")
+    source_type = DocumentSource(source_type_str) if source_type_str else None
+
     p_owner_names = all_metadata.get("primary_owners")
     s_owner_names = all_metadata.get("secondary_owners")
     p_owners = (
@@ -145,7 +149,7 @@ def _process_file(
             sections=[
                 Section(link=all_metadata.get("link"), text=file_content_raw.strip())
             ],
-            source=DocumentSource.FILE,
+            source=source_type or DocumentSource.FILE,
             semantic_identifier=file_display_name,
             title=title,
             doc_updated_at=final_time_updated,
@@ -175,7 +179,7 @@ class LocalFileConnector(LoadConnector):
 
     def load_from_state(self) -> GenerateDocumentsOutput:
         documents: list[Document] = []
-        token = current_tenant_id.set(self.tenant_id)
+        token = CURRENT_TENANT_ID_CONTEXTVAR.set(self.tenant_id)
 
         with get_session_with_tenant(self.tenant_id) as db_session:
             for file_path in self.file_locations:
@@ -199,7 +203,7 @@ class LocalFileConnector(LoadConnector):
             if documents:
                 yield documents
 
-        current_tenant_id.reset(token)
+        CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
 
 if __name__ == "__main__":

backend/danswer/connectors/fireflies/connector.py

@@ -0,0 +1,182 @@
+from collections.abc import Iterator
+from datetime import datetime
+from datetime import timezone
+from typing import List
+
+import requests
+
+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 LoadConnector
+from danswer.connectors.interfaces import PollConnector
+from danswer.connectors.interfaces import SecondsSinceUnixEpoch
+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.logger import setup_logger
+
+logger = setup_logger()
+
+_FIREFLIES_ID_PREFIX = "FIREFLIES_"
+
+_FIREFLIES_API_URL = "https://api.fireflies.ai/graphql"
+
+_FIREFLIES_TRANSCRIPT_QUERY_SIZE = 50  # Max page size is 50
+
+_FIREFLIES_API_QUERY = """
+    query Transcripts($fromDate: DateTime, $toDate: DateTime, $limit: Int!, $skip: Int!) {
+        transcripts(fromDate: $fromDate, toDate: $toDate, limit: $limit, skip: $skip) {
+            id
+            title
+            host_email
+            participants
+            date
+            transcript_url
+            sentences {
+                text
+                speaker_name
+            }
+        }
+    }
+"""
+
+
+def _create_doc_from_transcript(transcript: dict) -> Document | None:
+    meeting_text = ""
+    sentences = transcript.get("sentences", [])
+    if sentences:
+        for sentence in sentences:
+            meeting_text += sentence.get("speaker_name") or "Unknown Speaker"
+            meeting_text += ": " + sentence.get("text", "") + "\n\n"
+    else:
+        return None
+
+    meeting_link = transcript["transcript_url"]
+
+    fireflies_id = _FIREFLIES_ID_PREFIX + transcript["id"]
+
+    meeting_title = transcript["title"] or "No Title"
+
+    meeting_date_unix = transcript["date"]
+    meeting_date = datetime.fromtimestamp(meeting_date_unix / 1000, tz=timezone.utc)
+
+    meeting_host_email = transcript["host_email"]
+    host_email_user_info = [BasicExpertInfo(email=meeting_host_email)]
+
+    meeting_participants_email_list = []
+    for participant in transcript.get("participants", []):
+        if participant != meeting_host_email and participant:
+            meeting_participants_email_list.append(BasicExpertInfo(email=participant))
+
+    return Document(
+        id=fireflies_id,
+        sections=[
+            Section(
+                link=meeting_link,
+                text=meeting_text,
+            )
+        ],
+        source=DocumentSource.FIREFLIES,
+        semantic_identifier=meeting_title,
+        metadata={},
+        doc_updated_at=meeting_date,
+        primary_owners=host_email_user_info,
+        secondary_owners=meeting_participants_email_list,
+    )
+
+
+class FirefliesConnector(PollConnector, LoadConnector):
+    def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:
+        self.batch_size = batch_size
+
+    def load_credentials(self, credentials: dict[str, str]) -> None:
+        api_key = credentials.get("fireflies_api_key")
+
+        if not isinstance(api_key, str):
+            raise ConnectorMissingCredentialError(
+                "The Fireflies API key must be a string"
+            )
+
+        self.api_key = api_key
+
+        return None
+
+    def _fetch_transcripts(
+        self, start_datetime: str | None = None, end_datetime: str | None = None
+    ) -> Iterator[List[dict]]:
+        if self.api_key is None:
+            raise ConnectorMissingCredentialError("Missing API key")
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer " + self.api_key,
+        }
+
+        skip = 0
+        variables: dict[str, int | str] = {
+            "limit": _FIREFLIES_TRANSCRIPT_QUERY_SIZE,
+        }
+
+        if start_datetime:
+            variables["fromDate"] = start_datetime
+        if end_datetime:
+            variables["toDate"] = end_datetime
+
+        while True:
+            variables["skip"] = skip
+            response = requests.post(
+                _FIREFLIES_API_URL,
+                headers=headers,
+                json={"query": _FIREFLIES_API_QUERY, "variables": variables},
+            )
+
+            response.raise_for_status()
+
+            if response.status_code == 204:
+                break
+
+            recieved_transcripts = response.json()
+            parsed_transcripts = recieved_transcripts.get("data", {}).get(
+                "transcripts", []
+            )
+
+            yield parsed_transcripts
+
+            if len(parsed_transcripts) < _FIREFLIES_TRANSCRIPT_QUERY_SIZE:
+                break
+
+            skip += _FIREFLIES_TRANSCRIPT_QUERY_SIZE
+
+    def _process_transcripts(
+        self, start: str | None = None, end: str | None = None
+    ) -> GenerateDocumentsOutput:
+        doc_batch: List[Document] = []
+
+        for transcript_batch in self._fetch_transcripts(start, end):
+            for transcript in transcript_batch:
+                if doc := _create_doc_from_transcript(transcript):
+                    doc_batch.append(doc)
+
+                if len(doc_batch) >= self.batch_size:
+                    yield doc_batch
+                    doc_batch = []
+
+        if doc_batch:
+            yield doc_batch
+
+    def load_from_state(self) -> GenerateDocumentsOutput:
+        return self._process_transcripts()
+
+    def poll_source(
+        self, start_unixtime: SecondsSinceUnixEpoch, end_unixtime: SecondsSinceUnixEpoch
+    ) -> GenerateDocumentsOutput:
+        start_datetime = datetime.fromtimestamp(
+            start_unixtime, tz=timezone.utc
+        ).strftime("%Y-%m-%dT%H:%M:%S.000Z")
+        end_datetime = datetime.fromtimestamp(end_unixtime, tz=timezone.utc).strftime(
+            "%Y-%m-%dT%H:%M:%S.000Z"
+        )
+
+        yield from self._process_transcripts(start_datetime, end_datetime)

backend/danswer/connectors/freshdesk/connector.py

@@ -0,0 +1,239 @@
+import json
+from collections.abc import Iterator
+from datetime import datetime
+from datetime import timezone
+from typing import List
+
+import requests
+
+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 LoadConnector
+from danswer.connectors.interfaces import PollConnector
+from danswer.connectors.interfaces import SecondsSinceUnixEpoch
+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.logger import setup_logger
+
+logger = setup_logger()
+
+_FRESHDESK_ID_PREFIX = "FRESHDESK_"
+
+
+_TICKET_FIELDS_TO_INCLUDE = {
+    "fr_escalated",
+    "spam",
+    "priority",
+    "source",
+    "status",
+    "type",
+    "is_escalated",
+    "tags",
+    "nr_due_by",
+    "nr_escalated",
+    "cc_emails",
+    "fwd_emails",
+    "reply_cc_emails",
+    "ticket_cc_emails",
+    "support_email",
+    "to_emails",
+}
+
+_SOURCE_NUMBER_TYPE_MAP: dict[int, str] = {
+    1: "Email",
+    2: "Portal",
+    3: "Phone",
+    7: "Chat",
+    9: "Feedback Widget",
+    10: "Outbound Email",
+}
+
+_PRIORITY_NUMBER_TYPE_MAP: dict[int, str] = {
+    1: "low",
+    2: "medium",
+    3: "high",
+    4: "urgent",
+}
+
+_STATUS_NUMBER_TYPE_MAP: dict[int, str] = {
+    2: "open",
+    3: "pending",
+    4: "resolved",
+    5: "closed",
+}
+
+
+def _create_metadata_from_ticket(ticket: dict) -> dict:
+    metadata: dict[str, str | list[str]] = {}
+    # Combine all emails into a list so there are no repeated emails
+    email_data: set[str] = set()
+
+    for key, value in ticket.items():
+        # Skip fields that aren't useful for embedding
+        if key not in _TICKET_FIELDS_TO_INCLUDE:
+            continue
+
+        # Skip empty fields
+        if not value or value == "[]":
+            continue
+
+        # Convert strings or lists to strings
+        stringified_value: str | list[str]
+        if isinstance(value, list):
+            stringified_value = [str(item) for item in value]
+        else:
+            stringified_value = str(value)
+
+        if "email" in key:
+            if isinstance(stringified_value, list):
+                email_data.update(stringified_value)
+            else:
+                email_data.add(stringified_value)
+        else:
+            metadata[key] = stringified_value
+
+    if email_data:
+        metadata["emails"] = list(email_data)
+
+    # Convert source numbers to human-parsable string
+    if source_number := ticket.get("source"):
+        metadata["source"] = _SOURCE_NUMBER_TYPE_MAP.get(
+            source_number, "Unknown Source Type"
+        )
+
+    # Convert priority numbers to human-parsable string
+    if priority_number := ticket.get("priority"):
+        metadata["priority"] = _PRIORITY_NUMBER_TYPE_MAP.get(
+            priority_number, "Unknown Priority"
+        )
+
+    # Convert status to human-parsable string
+    if status_number := ticket.get("status"):
+        metadata["status"] = _STATUS_NUMBER_TYPE_MAP.get(
+            status_number, "Unknown Status"
+        )
+
+    due_by = datetime.fromisoformat(ticket["due_by"].replace("Z", "+00:00"))
+    metadata["overdue"] = str(datetime.now(timezone.utc) > due_by)
+
+    return metadata
+
+
+def _create_doc_from_ticket(ticket: dict, domain: str) -> Document:
+    # Use the ticket description as the text
+    text = f"Ticket description: {parse_html_page_basic(ticket.get('description_text', ''))}"
+    metadata = _create_metadata_from_ticket(ticket)
+
+    # This is also used in the ID because it is more unique than the just the ticket ID
+    link = f"https://{domain}.freshdesk.com/helpdesk/tickets/{ticket['id']}"
+
+    return Document(
+        id=_FRESHDESK_ID_PREFIX + link,
+        sections=[
+            Section(
+                link=link,
+                text=text,
+            )
+        ],
+        source=DocumentSource.FRESHDESK,
+        semantic_identifier=ticket["subject"],
+        metadata=metadata,
+        doc_updated_at=datetime.fromisoformat(
+            ticket["updated_at"].replace("Z", "+00:00")
+        ),
+    )
+
+
+class FreshdeskConnector(PollConnector, LoadConnector):
+    def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:
+        self.batch_size = batch_size
+
+    def load_credentials(self, credentials: dict[str, str | int]) -> None:
+        api_key = credentials.get("freshdesk_api_key")
+        domain = credentials.get("freshdesk_domain")
+        password = credentials.get("freshdesk_password")
+
+        if not all(isinstance(cred, str) for cred in [domain, api_key, password]):
+            raise ConnectorMissingCredentialError(
+                "All Freshdesk credentials must be strings"
+            )
+
+        self.api_key = str(api_key)
+        self.domain = str(domain)
+        self.password = str(password)
+
+    def _fetch_tickets(
+        self, start: datetime | None = None, end: datetime | None = None
+    ) -> Iterator[List[dict]]:
+        """
+        'end' is not currently used, so we may double fetch tickets created after the indexing
+        starts but before the actual call is made.
+
+        To use 'end' would require us to use the search endpoint but it has limitations,
+        namely having to fetch all IDs and then individually fetch each ticket because there is no
+        'include' field available for this endpoint:
+        https://developers.freshdesk.com/api/#filter_tickets
+        """
+        if self.api_key is None or self.domain is None or self.password is None:
+            raise ConnectorMissingCredentialError("freshdesk")
+
+        base_url = f"https://{self.domain}.freshdesk.com/api/v2/tickets"
+        params: dict[str, int | str] = {
+            "include": "description",
+            "per_page": 50,
+            "page": 1,
+        }
+
+        if start:
+            params["updated_since"] = start.isoformat()
+
+        while True:
+            response = requests.get(
+                base_url, auth=(self.api_key, self.password), params=params
+            )
+            response.raise_for_status()
+
+            if response.status_code == 204:
+                break
+
+            tickets = json.loads(response.content)
+            logger.info(
+                f"Fetched {len(tickets)} tickets from Freshdesk API (Page {params['page']})"
+            )
+
+            yield tickets
+
+            if len(tickets) < int(params["per_page"]):
+                break
+
+            params["page"] = int(params["page"]) + 1
+
+    def _process_tickets(
+        self, start: datetime | None = None, end: datetime | None = None
+    ) -> GenerateDocumentsOutput:
+        doc_batch: List[Document] = []
+
+        for ticket_batch in self._fetch_tickets(start, end):
+            for ticket in ticket_batch:
+                doc_batch.append(_create_doc_from_ticket(ticket, self.domain))
+
+                if len(doc_batch) >= self.batch_size:
+                    yield doc_batch
+                    doc_batch = []
+
+        if doc_batch:
+            yield doc_batch
+
+    def load_from_state(self) -> GenerateDocumentsOutput:
+        return self._process_tickets()
+
+    def poll_source(
+        self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch
+    ) -> GenerateDocumentsOutput:
+        start_datetime = datetime.fromtimestamp(start, tz=timezone.utc)
+        end_datetime = datetime.fromtimestamp(end, tz=timezone.utc)
+
+        yield from self._process_tickets(start_datetime, end_datetime)

backend/danswer/connectors/gitlab/connector.py

@@ -24,6 +24,9 @@ from danswer.connectors.models import Document
 from danswer.connectors.models import Section
 from danswer.utils.logger import setup_logger
 
+
+logger = setup_logger()
+
 # List of directories/Files to exclude
 exclude_patterns = [
     "logs",
@@ -31,7 +34,6 @@ exclude_patterns = [
     ".gitlab/",
     ".pre-commit-config.yaml",
 ]
-logger = setup_logger()
 
 
 def _batch_gitlab_objects(

backend/danswer/connectors/gmail/connector.py

@@ -1,221 +1,360 @@
 from base64 import urlsafe_b64decode
 from typing import Any
-from typing import cast
 from typing import Dict
 
 from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
 from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
-from googleapiclient import discovery  # type: ignore
 
 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.gmail.connector_auth import (
-    get_gmail_creds_for_authorized_user,
-)
-from danswer.connectors.gmail.connector_auth import (
-    get_gmail_creds_for_service_account,
-)
-from danswer.connectors.gmail.constants import (
-    DB_CREDENTIALS_DICT_DELEGATED_USER_KEY,
-)
-from danswer.connectors.gmail.constants import DB_CREDENTIALS_DICT_TOKEN_KEY
-from danswer.connectors.gmail.constants import (
-    GMAIL_DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+from danswer.connectors.google_utils.google_auth import get_google_creds
+from danswer.connectors.google_utils.google_utils import execute_paginated_retrieval
+from danswer.connectors.google_utils.resources import get_admin_service
+from danswer.connectors.google_utils.resources import get_gmail_service
+from danswer.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_PRIMARY_ADMIN_KEY,
 )
+from danswer.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR
+from danswer.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS
+from danswer.connectors.google_utils.shared_constants import SLIM_BATCH_SIZE
+from danswer.connectors.google_utils.shared_constants import USER_FIELDS
 from danswer.connectors.interfaces import GenerateDocumentsOutput
+from danswer.connectors.interfaces import GenerateSlimDocumentOutput
 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 Document
 from danswer.connectors.models import Section
+from danswer.connectors.models import SlimDocument
 from danswer.utils.logger import setup_logger
+from danswer.utils.retry_wrapper import retry_builder
+
 
 logger = setup_logger()
 
+# This is for the initial list call to get the thread ids
+THREAD_LIST_FIELDS = "nextPageToken, threads(id)"
 
-class GmailConnector(LoadConnector, PollConnector):
+# These are the fields to retrieve using the ID from the initial list call
+PARTS_FIELDS = "parts(body(data), mimeType)"
+PAYLOAD_FIELDS = f"payload(headers, {PARTS_FIELDS})"
+MESSAGES_FIELDS = f"messages(id, {PAYLOAD_FIELDS})"
+THREADS_FIELDS = f"threads(id, {MESSAGES_FIELDS})"
+THREAD_FIELDS = f"id, {MESSAGES_FIELDS}"
+
+EMAIL_FIELDS = [
+    "cc",
+    "bcc",
+    "from",
+    "to",
+]
+
+add_retries = retry_builder(tries=50, max_delay=30)
+
+
+def _build_time_range_query(
+    time_range_start: SecondsSinceUnixEpoch | None = None,
+    time_range_end: SecondsSinceUnixEpoch | None = None,
+) -> str | None:
+    query = ""
+    if time_range_start is not None and time_range_start != 0:
+        query += f"after:{int(time_range_start)}"
+    if time_range_end is not None and time_range_end != 0:
+        query += f" before:{int(time_range_end)}"
+    query = query.strip()
+
+    if len(query) == 0:
+        return None
+
+    return query
+
+
+def _clean_email_and_extract_name(email: str) -> tuple[str, str | None]:
+    email = email.strip()
+    if "<" in email and ">" in email:
+        # Handle format: "Display Name <email@domain.com>"
+        display_name = email[: email.find("<")].strip()
+        email_address = email[email.find("<") + 1 : email.find(">")].strip()
+        return email_address, display_name if display_name else None
+    else:
+        # Handle plain email address
+        return email.strip(), None
+
+
+def _get_owners_from_emails(emails: dict[str, str | None]) -> list[BasicExpertInfo]:
+    owners = []
+    for email, names in emails.items():
+        if names:
+            name_parts = names.split(" ")
+            first_name = " ".join(name_parts[:-1])
+            last_name = name_parts[-1]
+        else:
+            first_name = None
+            last_name = None
+        owners.append(
+            BasicExpertInfo(email=email, first_name=first_name, last_name=last_name)
+        )
+    return owners
+
+
+def _get_message_body(payload: dict[str, Any]) -> str:
+    parts = payload.get("parts", [])
+    message_body = ""
+    for part in parts:
+        mime_type = part.get("mimeType")
+        body = part.get("body")
+        if mime_type == "text/plain" and body:
+            data = body.get("data", "")
+            text = urlsafe_b64decode(data).decode()
+            message_body += text
+    return message_body
+
+
+def message_to_section(message: Dict[str, Any]) -> tuple[Section, dict[str, str]]:
+    link = f"https://mail.google.com/mail/u/0/#inbox/{message['id']}"
+
+    payload = message.get("payload", {})
+    headers = payload.get("headers", [])
+    metadata: dict[str, Any] = {}
+    for header in headers:
+        name = header.get("name").lower()
+        value = header.get("value")
+        if name in EMAIL_FIELDS:
+            metadata[name] = value
+        if name == "subject":
+            metadata["subject"] = value
+        if name == "date":
+            metadata["updated_at"] = value
+
+    if labels := message.get("labelIds"):
+        metadata["labels"] = labels
+
+    message_data = ""
+    for name, value in metadata.items():
+        # updated at isnt super useful for the llm
+        if name != "updated_at":
+            message_data += f"{name}: {value}\n"
+
+    message_body_text: str = _get_message_body(payload)
+
+    return Section(link=link, text=message_body_text + message_data), metadata
+
+
+def thread_to_document(full_thread: Dict[str, Any]) -> Document | None:
+    all_messages = full_thread.get("messages", [])
+    if not all_messages:
+        return None
+
+    sections = []
+    semantic_identifier = ""
+    updated_at = None
+    from_emails: dict[str, str | None] = {}
+    other_emails: dict[str, str | None] = {}
+    for message in all_messages:
+        section, message_metadata = message_to_section(message)
+        sections.append(section)
+
+        for name, value in message_metadata.items():
+            if name in EMAIL_FIELDS:
+                email, display_name = _clean_email_and_extract_name(value)
+                if name == "from":
+                    from_emails[email] = (
+                        display_name if not from_emails.get(email) else None
+                    )
+                else:
+                    other_emails[email] = (
+                        display_name if not other_emails.get(email) else None
+                    )
+
+        # If we haven't set the semantic identifier yet, set it to the subject of the first message
+        if not semantic_identifier:
+            semantic_identifier = message_metadata.get("subject", "")
+
+        if message_metadata.get("updated_at"):
+            updated_at = message_metadata.get("updated_at")
+
+    updated_at_datetime = None
+    if updated_at:
+        updated_at_datetime = time_str_to_utc(updated_at)
+
+    id = full_thread.get("id")
+    if not id:
+        raise ValueError("Thread ID is required")
+
+    primary_owners = _get_owners_from_emails(from_emails)
+    secondary_owners = _get_owners_from_emails(other_emails)
+
+    return Document(
+        id=id,
+        semantic_identifier=semantic_identifier,
+        sections=sections,
+        source=DocumentSource.GMAIL,
+        # This is used to perform permission sync
+        primary_owners=primary_owners,
+        secondary_owners=secondary_owners,
+        doc_updated_at=updated_at_datetime,
+        # Not adding emails to metadata because it's already in the sections
+        metadata={},
+    )
+
+
+class GmailConnector(LoadConnector, PollConnector, SlimConnector):
     def __init__(self, batch_size: int = INDEX_BATCH_SIZE) -> None:
         self.batch_size = batch_size
-        self.creds: OAuthCredentials | ServiceAccountCredentials | None = None
+
+        self._creds: OAuthCredentials | ServiceAccountCredentials | None = None
+        self._primary_admin_email: str | None = None
+
+    @property
+    def primary_admin_email(self) -> str:
+        if self._primary_admin_email is None:
+            raise RuntimeError(
+                "Primary admin email missing, "
+                "should not call this property "
+                "before calling load_credentials"
+            )
+        return self._primary_admin_email
+
+    @property
+    def google_domain(self) -> str:
+        if self._primary_admin_email is None:
+            raise RuntimeError(
+                "Primary admin email missing, "
+                "should not call this property "
+                "before calling load_credentials"
+            )
+        return self._primary_admin_email.split("@")[-1]
+
+    @property
+    def creds(self) -> OAuthCredentials | ServiceAccountCredentials:
+        if self._creds is None:
+            raise RuntimeError(
+                "Creds missing, "
+                "should not call this property "
+                "before calling load_credentials"
+            )
+        return self._creds
 
     def load_credentials(self, credentials: dict[str, Any]) -> dict[str, str] | None:
-        """Checks for two different types of credentials.
-        (1) A credential which holds a token acquired via a user going thorugh
-        the Google OAuth flow.
-        (2) A credential which holds a service account key JSON file, which
-        can then be used to impersonate any user in the workspace.
-        """
-        creds: OAuthCredentials | ServiceAccountCredentials | None = None
-        new_creds_dict = None
-        if DB_CREDENTIALS_DICT_TOKEN_KEY in credentials:
-            access_token_json_str = cast(
-                str, credentials[DB_CREDENTIALS_DICT_TOKEN_KEY]
-            )
-            creds = get_gmail_creds_for_authorized_user(
-                token_json_str=access_token_json_str
-            )
+        primary_admin_email = credentials[DB_CREDENTIALS_PRIMARY_ADMIN_KEY]
+        self._primary_admin_email = primary_admin_email
 
-            # tell caller to update token stored in DB if it has changed
-            # (e.g. the token has been refreshed)
-            new_creds_json_str = creds.to_json() if creds else ""
-            if new_creds_json_str != access_token_json_str:
-                new_creds_dict = {DB_CREDENTIALS_DICT_TOKEN_KEY: new_creds_json_str}
-
-        if GMAIL_DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY in credentials:
-            service_account_key_json_str = credentials[
-                GMAIL_DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY
-            ]
-            creds = get_gmail_creds_for_service_account(
-                service_account_key_json_str=service_account_key_json_str
-            )
-
-            # "Impersonate" a user if one is specified
-            delegated_user_email = cast(
-                str | None, credentials.get(DB_CREDENTIALS_DICT_DELEGATED_USER_KEY)
-            )
-            if delegated_user_email:
-                creds = creds.with_subject(delegated_user_email) if creds else None  # type: ignore
-
-        if creds is None:
-            raise PermissionError(
-                "Unable to access Gmail - unknown credential structure."
-            )
-
-        self.creds = creds
+        self._creds, new_creds_dict = get_google_creds(
+            credentials=credentials,
+            source=DocumentSource.GMAIL,
+        )
         return new_creds_dict
 
-    def _get_email_body(self, payload: dict[str, Any]) -> str:
-        parts = payload.get("parts", [])
-        email_body = ""
-        for part in parts:
-            mime_type = part.get("mimeType")
-            body = part.get("body")
-            if mime_type == "text/plain":
-                data = body.get("data", "")
-                text = urlsafe_b64decode(data).decode()
-                email_body += text
-        return email_body
+    def _get_all_user_emails(self) -> list[str]:
+        admin_service = get_admin_service(self.creds, self.primary_admin_email)
+        emails = []
+        for user in execute_paginated_retrieval(
+            retrieval_function=admin_service.users().list,
+            list_key="users",
+            fields=USER_FIELDS,
+            domain=self.google_domain,
+        ):
+            if email := user.get("primaryEmail"):
+                emails.append(email)
+        return emails
 
-    def _email_to_document(self, full_email: Dict[str, Any]) -> Document:
-        email_id = full_email["id"]
-        payload = full_email["payload"]
-        headers = payload.get("headers")
-        labels = full_email.get("labelIds", [])
-        metadata = {}
-        if headers:
-            for header in headers:
-                name = header.get("name").lower()
-                value = header.get("value")
-                if name in ["from", "to", "subject", "date", "cc", "bcc"]:
-                    metadata[name] = value
-        email_data = ""
-        for name, value in metadata.items():
-            email_data += f"{name}: {value}\n"
-        metadata["labels"] = labels
-        logger.debug(f"{email_data}")
-        email_body_text: str = self._get_email_body(payload)
-        date_str = metadata.get("date")
-        email_updated_at = time_str_to_utc(date_str) if date_str else None
-        link = f"https://mail.google.com/mail/u/0/#inbox/{email_id}"
-        return Document(
-            id=email_id,
-            sections=[Section(link=link, text=email_data + email_body_text)],
-            source=DocumentSource.GMAIL,
-            title=metadata.get("subject"),
-            semantic_identifier=metadata.get("subject", "Untitled Email"),
-            doc_updated_at=email_updated_at,
-            metadata=metadata,
-        )
-
-    @staticmethod
-    def _build_time_range_query(
-        time_range_start: SecondsSinceUnixEpoch | None = None,
-        time_range_end: SecondsSinceUnixEpoch | None = None,
-    ) -> str | None:
-        query = ""
-        if time_range_start is not None and time_range_start != 0:
-            query += f"after:{int(time_range_start)}"
-        if time_range_end is not None and time_range_end != 0:
-            query += f" before:{int(time_range_end)}"
-        query = query.strip()
-
-        if len(query) == 0:
-            return None
-
-        return query
-
-    def _fetch_mails_from_gmail(
+    def _fetch_threads(
         self,
         time_range_start: SecondsSinceUnixEpoch | None = None,
         time_range_end: SecondsSinceUnixEpoch | None = None,
     ) -> GenerateDocumentsOutput:
-        if self.creds is None:
-            raise PermissionError("Not logged into Gmail")
-        page_token = ""
-        query = GmailConnector._build_time_range_query(time_range_start, time_range_end)
-        service = discovery.build("gmail", "v1", credentials=self.creds)
-        while page_token is not None:
-            result = (
-                service.users()
-                .messages()
-                .list(
-                    userId="me",
-                    pageToken=page_token,
-                    q=query,
-                    maxResults=self.batch_size,
+        query = _build_time_range_query(time_range_start, time_range_end)
+        doc_batch = []
+        for user_email in self._get_all_user_emails():
+            gmail_service = get_gmail_service(self.creds, user_email)
+            for thread in execute_paginated_retrieval(
+                retrieval_function=gmail_service.users().threads().list,
+                list_key="threads",
+                userId=user_email,
+                fields=THREAD_LIST_FIELDS,
+                q=query,
+            ):
+                full_threads = execute_paginated_retrieval(
+                    retrieval_function=gmail_service.users().threads().get,
+                    list_key=None,
+                    userId=user_email,
+                    fields=THREAD_FIELDS,
+                    id=thread["id"],
                 )
-                .execute()
-            )
-            page_token = result.get("nextPageToken")
-            messages = result.get("messages", [])
-            doc_batch = []
-            for message in messages:
-                message_id = message["id"]
-                msg = (
-                    service.users()
-                    .messages()
-                    .get(userId="me", id=message_id, format="full")
-                    .execute()
-                )
-                doc = self._email_to_document(msg)
+                # full_threads is an iterator containing a single thread
+                # so we need to convert it to a list and grab the first element
+                full_thread = list(full_threads)[0]
+                doc = thread_to_document(full_thread)
+                if doc is None:
+                    continue
                 doc_batch.append(doc)
-            if len(doc_batch) > 0:
-                yield doc_batch
+                if len(doc_batch) > self.batch_size:
+                    yield doc_batch
+                    doc_batch = []
+        if doc_batch:
+            yield doc_batch
+
+    def _fetch_slim_threads(
+        self,
+        time_range_start: SecondsSinceUnixEpoch | None = None,
+        time_range_end: SecondsSinceUnixEpoch | None = None,
+    ) -> GenerateSlimDocumentOutput:
+        query = _build_time_range_query(time_range_start, time_range_end)
+        doc_batch = []
+        for user_email in self._get_all_user_emails():
+            gmail_service = get_gmail_service(self.creds, user_email)
+            for thread in execute_paginated_retrieval(
+                retrieval_function=gmail_service.users().threads().list,
+                list_key="threads",
+                userId=user_email,
+                fields=THREAD_LIST_FIELDS,
+                q=query,
+            ):
+                doc_batch.append(
+                    SlimDocument(
+                        id=thread["id"],
+                        perm_sync_data={"user_email": user_email},
+                    )
+                )
+                if len(doc_batch) > SLIM_BATCH_SIZE:
+                    yield doc_batch
+                    doc_batch = []
+        if doc_batch:
+            yield doc_batch
 
     def load_from_state(self) -> GenerateDocumentsOutput:
-        yield from self._fetch_mails_from_gmail()
+        try:
+            yield from self._fetch_threads()
+        except Exception as e:
+            if MISSING_SCOPES_ERROR_STR in str(e):
+                raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e
+            raise e
 
     def poll_source(
         self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch
     ) -> GenerateDocumentsOutput:
-        yield from self._fetch_mails_from_gmail(start, end)
+        try:
+            yield from self._fetch_threads(start, end)
+        except Exception as e:
+            if MISSING_SCOPES_ERROR_STR in str(e):
+                raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e
+            raise e
+
+    def retrieve_all_slim_documents(
+        self,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+    ) -> GenerateSlimDocumentOutput:
+        try:
+            yield from self._fetch_slim_threads(start, end)
+        except Exception as e:
+            if MISSING_SCOPES_ERROR_STR in str(e):
+                raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e
+            raise e
 
 
 if __name__ == "__main__":
-    import json
-    import os
-
-    service_account_json_path = os.environ.get("GOOGLE_SERVICE_ACCOUNT_KEY_JSON_PATH")
-    if not service_account_json_path:
-        raise ValueError(
-            "Please set GOOGLE_SERVICE_ACCOUNT_KEY_JSON_PATH environment variable"
-        )
-    with open(service_account_json_path) as f:
-        creds = json.load(f)
-
-    credentials_dict = {
-        DB_CREDENTIALS_DICT_TOKEN_KEY: json.dumps(creds),
-    }
-    delegated_user = os.environ.get("GMAIL_DELEGATED_USER")
-    if delegated_user:
-        credentials_dict[DB_CREDENTIALS_DICT_DELEGATED_USER_KEY] = delegated_user
-
-    connector = GmailConnector()
-    connector.load_credentials(
-        json.loads(credentials_dict[DB_CREDENTIALS_DICT_TOKEN_KEY])
-    )
-    document_batch_generator = connector.load_from_state()
-    for document_batch in document_batch_generator:
-        print(document_batch)
-        break
+    pass

backend/danswer/connectors/gmail/connector_auth.py

@@ -1,197 +0,0 @@
-import json
-from typing import cast
-from urllib.parse import parse_qs
-from urllib.parse import ParseResult
-from urllib.parse import urlparse
-
-from google.auth.transport.requests import Request  # type: ignore
-from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
-from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
-from google_auth_oauthlib.flow import InstalledAppFlow  # type: ignore
-from sqlalchemy.orm import Session
-
-from danswer.configs.app_configs import WEB_DOMAIN
-from danswer.configs.constants import DocumentSource
-from danswer.configs.constants import KV_CRED_KEY
-from danswer.configs.constants import KV_GMAIL_CRED_KEY
-from danswer.configs.constants import KV_GMAIL_SERVICE_ACCOUNT_KEY
-from danswer.connectors.gmail.constants import (
-    DB_CREDENTIALS_DICT_DELEGATED_USER_KEY,
-)
-from danswer.connectors.gmail.constants import DB_CREDENTIALS_DICT_TOKEN_KEY
-from danswer.connectors.gmail.constants import (
-    GMAIL_DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
-)
-from danswer.connectors.gmail.constants import SCOPES
-from danswer.db.credentials import update_credential_json
-from danswer.db.models import User
-from danswer.key_value_store.factory import get_kv_store
-from danswer.server.documents.models import CredentialBase
-from danswer.server.documents.models import GoogleAppCredentials
-from danswer.server.documents.models import GoogleServiceAccountKey
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def _build_frontend_gmail_redirect() -> str:
-    return f"{WEB_DOMAIN}/admin/connectors/gmail/auth/callback"
-
-
-def get_gmail_creds_for_authorized_user(
-    token_json_str: str,
-) -> OAuthCredentials | None:
-    creds_json = json.loads(token_json_str)
-    creds = OAuthCredentials.from_authorized_user_info(creds_json, SCOPES)
-    if creds.valid:
-        return creds
-
-    if creds.expired and creds.refresh_token:
-        try:
-            creds.refresh(Request())
-            if creds.valid:
-                logger.notice("Refreshed Gmail tokens.")
-                return creds
-        except Exception as e:
-            logger.exception(f"Failed to refresh gmail access token due to: {e}")
-            return None
-
-    return None
-
-
-def get_gmail_creds_for_service_account(
-    service_account_key_json_str: str,
-) -> ServiceAccountCredentials | None:
-    service_account_key = json.loads(service_account_key_json_str)
-    creds = ServiceAccountCredentials.from_service_account_info(
-        service_account_key, scopes=SCOPES
-    )
-    if not creds.valid or not creds.expired:
-        creds.refresh(Request())
-    return creds if creds.valid else None
-
-
-def verify_csrf(credential_id: int, state: str) -> None:
-    csrf = get_kv_store().load(KV_CRED_KEY.format(str(credential_id)))
-    if csrf != state:
-        raise PermissionError(
-            "State from Gmail Connector callback does not match expected"
-        )
-
-
-def get_gmail_auth_url(credential_id: int) -> str:
-    creds_str = str(get_kv_store().load(KV_GMAIL_CRED_KEY))
-    credential_json = json.loads(creds_str)
-    flow = InstalledAppFlow.from_client_config(
-        credential_json,
-        scopes=SCOPES,
-        redirect_uri=_build_frontend_gmail_redirect(),
-    )
-    auth_url, _ = flow.authorization_url(prompt="consent")
-
-    parsed_url = cast(ParseResult, urlparse(auth_url))
-    params = parse_qs(parsed_url.query)
-
-    get_kv_store().store(
-        KV_CRED_KEY.format(credential_id), params.get("state", [None])[0], encrypt=True
-    )  # type: ignore
-    return str(auth_url)
-
-
-def get_auth_url(credential_id: int) -> str:
-    creds_str = str(get_kv_store().load(KV_GMAIL_CRED_KEY))
-    credential_json = json.loads(creds_str)
-    flow = InstalledAppFlow.from_client_config(
-        credential_json,
-        scopes=SCOPES,
-        redirect_uri=_build_frontend_gmail_redirect(),
-    )
-    auth_url, _ = flow.authorization_url(prompt="consent")
-
-    parsed_url = cast(ParseResult, urlparse(auth_url))
-    params = parse_qs(parsed_url.query)
-
-    get_kv_store().store(
-        KV_CRED_KEY.format(credential_id), params.get("state", [None])[0], encrypt=True
-    )  # type: ignore
-    return str(auth_url)
-
-
-def update_gmail_credential_access_tokens(
-    auth_code: str,
-    credential_id: int,
-    user: User,
-    db_session: Session,
-) -> OAuthCredentials | None:
-    app_credentials = get_google_app_gmail_cred()
-    flow = InstalledAppFlow.from_client_config(
-        app_credentials.model_dump(),
-        scopes=SCOPES,
-        redirect_uri=_build_frontend_gmail_redirect(),
-    )
-    flow.fetch_token(code=auth_code)
-    creds = flow.credentials
-    token_json_str = creds.to_json()
-    new_creds_dict = {DB_CREDENTIALS_DICT_TOKEN_KEY: token_json_str}
-
-    if not update_credential_json(credential_id, new_creds_dict, user, db_session):
-        return None
-    return creds
-
-
-def build_service_account_creds(
-    delegated_user_email: str | None = None,
-) -> CredentialBase:
-    service_account_key = get_gmail_service_account_key()
-
-    credential_dict = {
-        GMAIL_DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: service_account_key.json(),
-    }
-    if delegated_user_email:
-        credential_dict[DB_CREDENTIALS_DICT_DELEGATED_USER_KEY] = delegated_user_email
-
-    return CredentialBase(
-        source=DocumentSource.GMAIL,
-        credential_json=credential_dict,
-        admin_public=True,
-    )
-
-
-def get_google_app_gmail_cred() -> GoogleAppCredentials:
-    creds_str = str(get_kv_store().load(KV_GMAIL_CRED_KEY))
-    return GoogleAppCredentials(**json.loads(creds_str))
-
-
-def upsert_google_app_gmail_cred(app_credentials: GoogleAppCredentials) -> None:
-    get_kv_store().store(KV_GMAIL_CRED_KEY, app_credentials.json(), encrypt=True)
-
-
-def delete_google_app_gmail_cred() -> None:
-    get_kv_store().delete(KV_GMAIL_CRED_KEY)
-
-
-def get_gmail_service_account_key() -> GoogleServiceAccountKey:
-    creds_str = str(get_kv_store().load(KV_GMAIL_SERVICE_ACCOUNT_KEY))
-    return GoogleServiceAccountKey(**json.loads(creds_str))
-
-
-def upsert_gmail_service_account_key(
-    service_account_key: GoogleServiceAccountKey,
-) -> None:
-    get_kv_store().store(
-        KV_GMAIL_SERVICE_ACCOUNT_KEY, service_account_key.json(), encrypt=True
-    )
-
-
-def upsert_service_account_key(service_account_key: GoogleServiceAccountKey) -> None:
-    get_kv_store().store(
-        KV_GMAIL_SERVICE_ACCOUNT_KEY, service_account_key.json(), encrypt=True
-    )
-
-
-def delete_gmail_service_account_key() -> None:
-    get_kv_store().delete(KV_GMAIL_SERVICE_ACCOUNT_KEY)
-
-
-def delete_service_account_key() -> None:
-    get_kv_store().delete(KV_GMAIL_SERVICE_ACCOUNT_KEY)

backend/danswer/connectors/gmail/constants.py

@@ -1,4 +0,0 @@
-DB_CREDENTIALS_DICT_TOKEN_KEY = "gmail_tokens"
-GMAIL_DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY = "gmail_service_account_key"
-DB_CREDENTIALS_DICT_DELEGATED_USER_KEY = "gmail_delegated_user"
-SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]

backend/danswer/connectors/google_drive/connector.py

@@ -1,556 +1,400 @@
-import io
+from collections.abc import Callable
 from collections.abc import Iterator
-from collections.abc import Sequence
-from datetime import datetime
-from datetime import timezone
-from enum import Enum
-from itertools import chain
+from concurrent.futures import as_completed
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 from typing import Any
 
 from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
 from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
-from googleapiclient import discovery  # type: ignore
-from googleapiclient.errors import HttpError  # type: ignore
 
-from danswer.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE
-from danswer.configs.app_configs import GOOGLE_DRIVE_FOLLOW_SHORTCUTS
-from danswer.configs.app_configs import GOOGLE_DRIVE_INCLUDE_SHARED
-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,
+from danswer.connectors.google_drive.doc_conversion import build_slim_document
+from danswer.connectors.google_drive.doc_conversion import (
+    convert_drive_item_to_document,
 )
-from danswer.connectors.google_drive.constants import (
-    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+from danswer.connectors.google_drive.file_retrieval import crawl_folders_for_files
+from danswer.connectors.google_drive.file_retrieval import get_all_files_in_my_drive
+from danswer.connectors.google_drive.file_retrieval import get_files_in_shared_drive
+from danswer.connectors.google_drive.models import GoogleDriveFileType
+from danswer.connectors.google_utils.google_auth import get_google_creds
+from danswer.connectors.google_utils.google_utils import execute_paginated_retrieval
+from danswer.connectors.google_utils.resources import get_admin_service
+from danswer.connectors.google_utils.resources import get_drive_service
+from danswer.connectors.google_utils.resources import get_google_docs_service
+from danswer.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_PRIMARY_ADMIN_KEY,
 )
+from danswer.connectors.google_utils.shared_constants import MISSING_SCOPES_ERROR_STR
+from danswer.connectors.google_utils.shared_constants import ONYX_SCOPE_INSTRUCTIONS
+from danswer.connectors.google_utils.shared_constants import SCOPE_DOC_URL
+from danswer.connectors.google_utils.shared_constants import SLIM_BATCH_SIZE
+from danswer.connectors.google_utils.shared_constants import USER_FIELDS
 from danswer.connectors.interfaces import GenerateDocumentsOutput
+from danswer.connectors.interfaces import GenerateSlimDocumentOutput
 from danswer.connectors.interfaces import LoadConnector
 from danswer.connectors.interfaces import PollConnector
 from danswer.connectors.interfaces import SecondsSinceUnixEpoch
-from danswer.connectors.models import Document
-from danswer.connectors.models import Section
-from danswer.file_processing.extract_file_text import docx_to_text
-from danswer.file_processing.extract_file_text import pptx_to_text
-from danswer.file_processing.extract_file_text import read_pdf_file
-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.connectors.interfaces import SlimConnector
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
-
-DRIVE_FOLDER_TYPE = "application/vnd.google-apps.folder"
-DRIVE_SHORTCUT_TYPE = "application/vnd.google-apps.shortcut"
-UNSUPPORTED_FILE_TYPE_CONTENT = ""  # keep empty for now
+# TODO: Improve this by using the batch utility: https://googleapis.github.io/google-api-python-client/docs/batch.html
+# All file retrievals could be batched and made at once
 
 
-class GDriveMimeType(str, Enum):
-    DOC = "application/vnd.google-apps.document"
-    SPREADSHEET = "application/vnd.google-apps.spreadsheet"
-    PDF = "application/pdf"
-    WORD_DOC = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
-    PPT = "application/vnd.google-apps.presentation"
-    POWERPOINT = (
-        "application/vnd.openxmlformats-officedocument.presentationml.presentation"
-    )
-    PLAIN_TEXT = "text/plain"
-    MARKDOWN = "text/markdown"
+def _extract_str_list_from_comma_str(string: str | None) -> list[str]:
+    if not string:
+        return []
+    return [s.strip() for s in string.split(",") if s.strip()]
 
 
-GoogleDriveFileType = dict[str, Any]
-
-# Google Drive APIs are quite flakey and may 500 for an
-# extended period of time. Trying to combat here by adding a very
-# long retry period (~20 minutes of trying every minute)
-add_retries = retry_builder(tries=50, max_delay=30)
+def _extract_ids_from_urls(urls: list[str]) -> list[str]:
+    return [url.split("/")[-1] for url in urls]
 
 
-def _run_drive_file_query(
-    service: discovery.Resource,
-    query: str,
-    continue_on_failure: bool,
-    include_shared: bool = GOOGLE_DRIVE_INCLUDE_SHARED,
-    follow_shortcuts: bool = GOOGLE_DRIVE_FOLLOW_SHORTCUTS,
-    batch_size: int = INDEX_BATCH_SIZE,
-) -> Iterator[GoogleDriveFileType]:
-    next_page_token = ""
-    while next_page_token is not None:
-        logger.debug(f"Running Google Drive fetch with query: {query}")
-        results = add_retries(
-            lambda: (
-                service.files()
-                .list(
-                    corpora="allDrives"
-                    if include_shared
-                    else "user",  # needed to search through shared drives
-                    pageSize=batch_size,
-                    supportsAllDrives=include_shared,
-                    includeItemsFromAllDrives=include_shared,
-                    fields=(
-                        "nextPageToken, files(mimeType, id, name, permissions, "
-                        "modifiedTime, webViewLink, shortcutDetails)"
-                    ),
-                    pageToken=next_page_token,
-                    q=query,
-                )
-                .execute()
-            )
-        )()
-        next_page_token = results.get("nextPageToken")
-        files = results["files"]
-        for file in files:
-            if follow_shortcuts and "shortcutDetails" in file:
-                try:
-                    file_shortcut_points_to = add_retries(
-                        lambda: (
-                            service.files()
-                            .get(
-                                fileId=file["shortcutDetails"]["targetId"],
-                                supportsAllDrives=include_shared,
-                                fields="mimeType, id, name, modifiedTime, webViewLink, permissions, shortcutDetails",
-                            )
-                            .execute()
-                        )
-                    )()
-                    yield file_shortcut_points_to
-                except HttpError:
-                    logger.error(
-                        f"Failed to follow shortcut with details: {file['shortcutDetails']}"
-                    )
-                    if continue_on_failure:
-                        continue
-                    raise
-            else:
-                yield file
-
-
-def _get_folder_id(
-    service: discovery.Resource,
-    parent_id: str,
-    folder_name: str,
-    include_shared: bool,
-    follow_shortcuts: bool,
-) -> str | None:
-    """
-    Get the ID of a folder given its name and the ID of its parent folder.
-    """
-    query = f"'{parent_id}' in parents and name='{folder_name}' and "
-    if follow_shortcuts:
-        query += f"(mimeType='{DRIVE_FOLDER_TYPE}' or mimeType='{DRIVE_SHORTCUT_TYPE}')"
-    else:
-        query += f"mimeType='{DRIVE_FOLDER_TYPE}'"
-
-    # TODO: support specifying folder path in shared drive rather than just `My Drive`
-    results = add_retries(
-        lambda: (
-            service.files()
-            .list(
-                q=query,
-                spaces="drive",
-                fields="nextPageToken, files(id, name, shortcutDetails)",
-                supportsAllDrives=include_shared,
-                includeItemsFromAllDrives=include_shared,
-            )
-            .execute()
-        )
-    )()
-    items = results.get("files", [])
-
-    folder_id = None
-    if items:
-        if follow_shortcuts and "shortcutDetails" in items[0]:
-            folder_id = items[0]["shortcutDetails"]["targetId"]
-        else:
-            folder_id = items[0]["id"]
-    return folder_id
-
-
-def _get_folders(
-    service: discovery.Resource,
-    continue_on_failure: bool,
-    folder_id: str | None = None,  # if specified, only fetches files within this folder
-    include_shared: bool = GOOGLE_DRIVE_INCLUDE_SHARED,
-    follow_shortcuts: bool = GOOGLE_DRIVE_FOLLOW_SHORTCUTS,
-    batch_size: int = INDEX_BATCH_SIZE,
-) -> Iterator[GoogleDriveFileType]:
-    query = f"mimeType = '{DRIVE_FOLDER_TYPE}' "
-    if follow_shortcuts:
-        query = "(" + query + f" or mimeType = '{DRIVE_SHORTCUT_TYPE}'" + ") "
-
-    if folder_id:
-        query += f"and '{folder_id}' in parents "
-    query = query.rstrip()  # remove the trailing space(s)
-
-    for file in _run_drive_file_query(
-        service=service,
-        query=query,
-        continue_on_failure=continue_on_failure,
-        include_shared=include_shared,
-        follow_shortcuts=follow_shortcuts,
-        batch_size=batch_size,
-    ):
-        # Need to check this since file may have been a target of a shortcut
-        # and not necessarily a folder
-        if file["mimeType"] == DRIVE_FOLDER_TYPE:
-            yield file
-        else:
-            pass
-
-
-def _get_files(
-    service: discovery.Resource,
-    continue_on_failure: bool,
-    time_range_start: SecondsSinceUnixEpoch | None = None,
-    time_range_end: SecondsSinceUnixEpoch | None = None,
-    folder_id: str | None = None,  # if specified, only fetches files within this folder
-    include_shared: bool = GOOGLE_DRIVE_INCLUDE_SHARED,
-    follow_shortcuts: bool = GOOGLE_DRIVE_FOLLOW_SHORTCUTS,
-    batch_size: int = INDEX_BATCH_SIZE,
-) -> Iterator[GoogleDriveFileType]:
-    query = f"mimeType != '{DRIVE_FOLDER_TYPE}' "
-    if time_range_start is not None:
-        time_start = datetime.utcfromtimestamp(time_range_start).isoformat() + "Z"
-        query += f"and modifiedTime >= '{time_start}' "
-    if time_range_end is not None:
-        time_stop = datetime.utcfromtimestamp(time_range_end).isoformat() + "Z"
-        query += f"and modifiedTime <= '{time_stop}' "
-    if folder_id:
-        query += f"and '{folder_id}' in parents "
-    query = query.rstrip()  # remove the trailing space(s)
-
-    files = _run_drive_file_query(
-        service=service,
-        query=query,
-        continue_on_failure=continue_on_failure,
-        include_shared=include_shared,
-        follow_shortcuts=follow_shortcuts,
-        batch_size=batch_size,
+def _convert_single_file(
+    creds: Any, primary_admin_email: str, file: dict[str, Any]
+) -> Any:
+    user_email = file.get("owners", [{}])[0].get("emailAddress") or primary_admin_email
+    user_drive_service = get_drive_service(creds, user_email=user_email)
+    docs_service = get_google_docs_service(creds, user_email=user_email)
+    return convert_drive_item_to_document(
+        file=file,
+        drive_service=user_drive_service,
+        docs_service=docs_service,
     )
 
-    return files
+
+def _process_files_batch(
+    files: list[GoogleDriveFileType], convert_func: Callable, batch_size: int
+) -> GenerateDocumentsOutput:
+    doc_batch = []
+    with ThreadPoolExecutor(max_workers=min(16, len(files))) as executor:
+        for doc in executor.map(convert_func, files):
+            if doc:
+                doc_batch.append(doc)
+                if len(doc_batch) >= batch_size:
+                    yield doc_batch
+                    doc_batch = []
+    if doc_batch:
+        yield doc_batch
 
 
-def get_all_files_batched(
-    service: discovery.Resource,
-    continue_on_failure: bool,
-    include_shared: bool = GOOGLE_DRIVE_INCLUDE_SHARED,
-    follow_shortcuts: bool = GOOGLE_DRIVE_FOLLOW_SHORTCUTS,
-    batch_size: int = INDEX_BATCH_SIZE,
-    time_range_start: SecondsSinceUnixEpoch | None = None,
-    time_range_end: SecondsSinceUnixEpoch | None = None,
-    folder_id: str | None = None,  # if specified, only fetches files within this folder
-    # if True, will fetch files in sub-folders of the specified folder ID.
-    # Only applies if folder_id is specified.
-    traverse_subfolders: bool = True,
-    folder_ids_traversed: list[str] | None = None,
-) -> Iterator[list[GoogleDriveFileType]]:
-    """Gets all files matching the criteria specified by the args from Google Drive
-    in batches of size `batch_size`.
-    """
-    found_files = _get_files(
-        service=service,
-        continue_on_failure=continue_on_failure,
-        time_range_start=time_range_start,
-        time_range_end=time_range_end,
-        folder_id=folder_id,
-        include_shared=include_shared,
-        follow_shortcuts=follow_shortcuts,
-        batch_size=batch_size,
-    )
-    yield from batch_generator(
-        items=found_files,
-        batch_size=batch_size,
-        pre_batch_yield=lambda batch_files: logger.debug(
-            f"Parseable Documents in batch: {[file['name'] for file in batch_files]}"
-        ),
-    )
-
-    if traverse_subfolders and folder_id is not None:
-        folder_ids_traversed = folder_ids_traversed or []
-        subfolders = _get_folders(
-            service=service,
-            folder_id=folder_id,
-            continue_on_failure=continue_on_failure,
-            include_shared=include_shared,
-            follow_shortcuts=follow_shortcuts,
-            batch_size=batch_size,
-        )
-        for subfolder in subfolders:
-            if subfolder["id"] not in folder_ids_traversed:
-                logger.info("Fetching all files in subfolder: " + subfolder["name"])
-                folder_ids_traversed.append(subfolder["id"])
-                yield from get_all_files_batched(
-                    service=service,
-                    continue_on_failure=continue_on_failure,
-                    include_shared=include_shared,
-                    follow_shortcuts=follow_shortcuts,
-                    batch_size=batch_size,
-                    time_range_start=time_range_start,
-                    time_range_end=time_range_end,
-                    folder_id=subfolder["id"],
-                    traverse_subfolders=traverse_subfolders,
-                    folder_ids_traversed=folder_ids_traversed,
-                )
-            else:
-                logger.debug(
-                    "Skipping subfolder since already traversed: " + subfolder["name"]
-                )
-
-
-def extract_text(file: dict[str, str], service: discovery.Resource) -> str:
-    mime_type = file["mimeType"]
-
-    if mime_type not in set(item.value for item in GDriveMimeType):
-        # Unsupported file types can still have a title, finding this way is still useful
-        return UNSUPPORTED_FILE_TYPE_CONTENT
-
-    if mime_type in [
-        GDriveMimeType.DOC.value,
-        GDriveMimeType.PPT.value,
-        GDriveMimeType.SPREADSHEET.value,
-    ]:
-        export_mime_type = (
-            "text/plain"
-            if mime_type != GDriveMimeType.SPREADSHEET.value
-            else "text/csv"
-        )
-        return (
-            service.files()
-            .export(fileId=file["id"], mimeType=export_mime_type)
-            .execute()
-            .decode("utf-8")
-        )
-    elif mime_type in [
-        GDriveMimeType.PLAIN_TEXT.value,
-        GDriveMimeType.MARKDOWN.value,
-    ]:
-        return service.files().get_media(fileId=file["id"]).execute().decode("utf-8")
-    if mime_type in [
-        GDriveMimeType.WORD_DOC.value,
-        GDriveMimeType.POWERPOINT.value,
-        GDriveMimeType.PDF.value,
-    ]:
-        response = service.files().get_media(fileId=file["id"]).execute()
-        if get_unstructured_api_key():
-            return unstructured_to_text(
-                file=io.BytesIO(response), file_name=file.get("name", file["id"])
-            )
-
-        if mime_type == GDriveMimeType.WORD_DOC.value:
-            return docx_to_text(file=io.BytesIO(response))
-        elif mime_type == GDriveMimeType.PDF.value:
-            text, _ = read_pdf_file(file=io.BytesIO(response))
-            return text
-        elif mime_type == GDriveMimeType.POWERPOINT.value:
-            return pptx_to_text(file=io.BytesIO(response))
-
-    return UNSUPPORTED_FILE_TYPE_CONTENT
-
-
-class GoogleDriveConnector(LoadConnector, PollConnector):
+class GoogleDriveConnector(LoadConnector, PollConnector, SlimConnector):
     def __init__(
         self,
-        # optional list of folder paths e.g. "[My Folder/My Subfolder]"
-        # if specified, will only index files in these folders
-        folder_paths: list[str] | None = None,
+        include_shared_drives: bool = True,
+        shared_drive_urls: str | None = None,
+        include_my_drives: bool = True,
+        my_drive_emails: str | None = None,
+        shared_folder_urls: str | None = None,
         batch_size: int = INDEX_BATCH_SIZE,
-        include_shared: bool = GOOGLE_DRIVE_INCLUDE_SHARED,
-        follow_shortcuts: bool = GOOGLE_DRIVE_FOLLOW_SHORTCUTS,
-        only_org_public: bool = GOOGLE_DRIVE_ONLY_ORG_PUBLIC,
-        continue_on_failure: bool = CONTINUE_ON_CONNECTOR_FAILURE,
+        # OLD PARAMETERS
+        folder_paths: list[str] | None = None,
+        include_shared: bool | None = None,
+        follow_shortcuts: bool | None = None,
+        only_org_public: bool | None = None,
+        continue_on_failure: bool | None = None,
     ) -> None:
-        self.folder_paths = folder_paths or []
+        # Check for old input parameters
+        if (
+            folder_paths is not None
+            or include_shared is not None
+            or follow_shortcuts is not None
+            or only_org_public is not None
+            or continue_on_failure is not None
+        ):
+            logger.exception(
+                "Google Drive connector received old input parameters. "
+                "Please visit the docs for help with the new setup: "
+                f"{SCOPE_DOC_URL}"
+            )
+            raise ValueError(
+                "Google Drive connector received old input parameters. "
+                "Please visit the docs for help with the new setup: "
+                f"{SCOPE_DOC_URL}"
+            )
+
+        if (
+            not include_shared_drives
+            and not include_my_drives
+            and not shared_folder_urls
+        ):
+            raise ValueError(
+                "At least one of include_shared_drives, include_my_drives,"
+                " or shared_folder_urls must be true"
+            )
+
         self.batch_size = batch_size
-        self.include_shared = include_shared
-        self.follow_shortcuts = follow_shortcuts
-        self.only_org_public = only_org_public
-        self.continue_on_failure = continue_on_failure
-        self.creds: OAuthCredentials | ServiceAccountCredentials | None = None
 
-    @staticmethod
-    def _process_folder_paths(
-        service: discovery.Resource,
-        folder_paths: list[str],
-        include_shared: bool,
-        follow_shortcuts: bool,
-    ) -> list[str]:
-        """['Folder/Sub Folder'] -> ['<FOLDER_ID>']"""
-        folder_ids: list[str] = []
-        for path in folder_paths:
-            folder_names = path.split("/")
-            parent_id = "root"
-            for folder_name in folder_names:
-                found_parent_id = _get_folder_id(
-                    service=service,
-                    parent_id=parent_id,
-                    folder_name=folder_name,
-                    include_shared=include_shared,
-                    follow_shortcuts=follow_shortcuts,
-                )
-                if found_parent_id is None:
-                    raise ValueError(
-                        (
-                            f"Folder '{folder_name}' in path '{path}' "
-                            "not found in Google Drive"
-                        )
-                    )
-                parent_id = found_parent_id
-            folder_ids.append(parent_id)
+        self.include_shared_drives = include_shared_drives
+        shared_drive_url_list = _extract_str_list_from_comma_str(shared_drive_urls)
+        self._requested_shared_drive_ids = set(
+            _extract_ids_from_urls(shared_drive_url_list)
+        )
 
-        return folder_ids
+        self.include_my_drives = include_my_drives
+        self._requested_my_drive_emails = set(
+            _extract_str_list_from_comma_str(my_drive_emails)
+        )
+
+        shared_folder_url_list = _extract_str_list_from_comma_str(shared_folder_urls)
+        self._requested_folder_ids = set(_extract_ids_from_urls(shared_folder_url_list))
+
+        self._primary_admin_email: str | None = None
+
+        self._creds: OAuthCredentials | ServiceAccountCredentials | None = None
+
+        self._retrieved_ids: set[str] = set()
+
+    @property
+    def primary_admin_email(self) -> str:
+        if self._primary_admin_email is None:
+            raise RuntimeError(
+                "Primary admin email missing, "
+                "should not call this property "
+                "before calling load_credentials"
+            )
+        return self._primary_admin_email
+
+    @property
+    def google_domain(self) -> str:
+        if self._primary_admin_email is None:
+            raise RuntimeError(
+                "Primary admin email missing, "
+                "should not call this property "
+                "before calling load_credentials"
+            )
+        return self._primary_admin_email.split("@")[-1]
+
+    @property
+    def creds(self) -> OAuthCredentials | ServiceAccountCredentials:
+        if self._creds is None:
+            raise RuntimeError(
+                "Creds missing, "
+                "should not call this property "
+                "before calling load_credentials"
+            )
+        return self._creds
 
     def load_credentials(self, credentials: dict[str, Any]) -> dict[str, str] | None:
-        """Checks for two different types of credentials.
-        (1) A credential which holds a token acquired via a user going thorough
-        the Google OAuth flow.
-        (2) A credential which holds a service account key JSON file, which
-        can then be used to impersonate any user in the workspace.
-        """
-        creds, new_creds_dict = get_google_drive_creds(credentials)
-        self.creds = creds
+        primary_admin_email = credentials[DB_CREDENTIALS_PRIMARY_ADMIN_KEY]
+        self._primary_admin_email = primary_admin_email
+
+        self._creds, new_creds_dict = get_google_creds(
+            credentials=credentials,
+            source=DocumentSource.GOOGLE_DRIVE,
+        )
         return new_creds_dict
 
-    def _fetch_docs_from_drive(
+    def _update_traversed_parent_ids(self, folder_id: str) -> None:
+        self._retrieved_ids.add(folder_id)
+
+    def _get_all_user_emails(self, admins_only: bool) -> list[str]:
+        admin_service = get_admin_service(
+            creds=self.creds,
+            user_email=self.primary_admin_email,
+        )
+        query = "isAdmin=true" if admins_only else "isAdmin=false"
+        emails = []
+        for user in execute_paginated_retrieval(
+            retrieval_function=admin_service.users().list,
+            list_key="users",
+            fields=USER_FIELDS,
+            domain=self.google_domain,
+            query=query,
+        ):
+            if email := user.get("primaryEmail"):
+                emails.append(email)
+        return emails
+
+    def _get_all_drive_ids(self) -> set[str]:
+        primary_drive_service = get_drive_service(
+            creds=self.creds,
+            user_email=self.primary_admin_email,
+        )
+        all_drive_ids = set()
+        for drive in execute_paginated_retrieval(
+            retrieval_function=primary_drive_service.drives().list,
+            list_key="drives",
+            useDomainAdminAccess=True,
+            fields="drives(id)",
+        ):
+            all_drive_ids.add(drive["id"])
+        return all_drive_ids
+
+    def _initialize_all_class_variables(self) -> None:
+        # Get all user emails
+        # Get admins first becuase they are more likely to have access to the most files
+        user_emails = [self.primary_admin_email]
+        for admins_only in [True, False]:
+            for email in self._get_all_user_emails(admins_only=admins_only):
+                if email not in user_emails:
+                    user_emails.append(email)
+        self._all_org_emails = user_emails
+
+        self._all_drive_ids: set[str] = self._get_all_drive_ids()
+
+        # remove drive ids from the folder ids because they are queried differently
+        self._requested_folder_ids -= self._all_drive_ids
+
+        # Remove drive_ids that are not in the all_drive_ids and check them as folders instead
+        invalid_drive_ids = self._requested_shared_drive_ids - self._all_drive_ids
+        if invalid_drive_ids:
+            logger.warning(
+                f"Some shared drive IDs were not found. IDs: {invalid_drive_ids}"
+            )
+            logger.warning("Checking for folder access instead...")
+            self._requested_folder_ids.update(invalid_drive_ids)
+
+        if not self.include_shared_drives:
+            self._requested_shared_drive_ids = set()
+        elif not self._requested_shared_drive_ids:
+            self._requested_shared_drive_ids = self._all_drive_ids
+
+    def _impersonate_user_for_retrieval(
+        self,
+        user_email: str,
+        is_slim: bool,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+    ) -> Iterator[GoogleDriveFileType]:
+        drive_service = get_drive_service(self.creds, user_email)
+        if self.include_my_drives and (
+            not self._requested_my_drive_emails
+            or user_email in self._requested_my_drive_emails
+        ):
+            yield from get_all_files_in_my_drive(
+                service=drive_service,
+                update_traversed_ids_func=self._update_traversed_parent_ids,
+                is_slim=is_slim,
+                start=start,
+                end=end,
+            )
+
+        remaining_drive_ids = self._requested_shared_drive_ids - self._retrieved_ids
+        for drive_id in remaining_drive_ids:
+            yield from get_files_in_shared_drive(
+                service=drive_service,
+                drive_id=drive_id,
+                is_slim=is_slim,
+                update_traversed_ids_func=self._update_traversed_parent_ids,
+                start=start,
+                end=end,
+            )
+
+        remaining_folders = self._requested_folder_ids - self._retrieved_ids
+        for folder_id in remaining_folders:
+            yield from crawl_folders_for_files(
+                service=drive_service,
+                parent_id=folder_id,
+                traversed_parent_ids=self._retrieved_ids,
+                update_traversed_ids_func=self._update_traversed_parent_ids,
+                start=start,
+                end=end,
+            )
+
+    def _fetch_drive_items(
+        self,
+        is_slim: bool,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+    ) -> Iterator[GoogleDriveFileType]:
+        self._initialize_all_class_variables()
+
+        # Process users in parallel using ThreadPoolExecutor
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            future_to_email = {
+                executor.submit(
+                    self._impersonate_user_for_retrieval, email, is_slim, start, end
+                ): email
+                for email in self._all_org_emails
+            }
+
+            # Yield results as they complete
+            for future in as_completed(future_to_email):
+                yield from future.result()
+
+        remaining_folders = self._requested_folder_ids - self._retrieved_ids
+        if remaining_folders:
+            logger.warning(
+                f"Some folders/drives were not retrieved. IDs: {remaining_folders}"
+            )
+
+    def _extract_docs_from_google_drive(
         self,
         start: SecondsSinceUnixEpoch | None = None,
         end: SecondsSinceUnixEpoch | None = None,
     ) -> GenerateDocumentsOutput:
-        if self.creds is None:
-            raise PermissionError("Not logged into Google Drive")
-
-        service = discovery.build("drive", "v3", credentials=self.creds)
-        folder_ids: Sequence[str | None] = self._process_folder_paths(
-            service, self.folder_paths, self.include_shared, self.follow_shortcuts
+        # Create a larger process pool for file conversion
+        convert_func = partial(
+            _convert_single_file, self.creds, self.primary_admin_email
         )
-        if not folder_ids:
-            folder_ids = [None]
 
-        file_batches = chain(
-            *[
-                get_all_files_batched(
-                    service=service,
-                    continue_on_failure=self.continue_on_failure,
-                    include_shared=self.include_shared,
-                    follow_shortcuts=self.follow_shortcuts,
-                    batch_size=self.batch_size,
-                    time_range_start=start,
-                    time_range_end=end,
-                    folder_id=folder_id,
-                    traverse_subfolders=True,
+        # Process files in larger batches
+        LARGE_BATCH_SIZE = self.batch_size * 4
+        files_to_process = []
+        # Gather the files into batches to be processed in parallel
+        for file in self._fetch_drive_items(is_slim=False, start=start, end=end):
+            files_to_process.append(file)
+            if len(files_to_process) >= LARGE_BATCH_SIZE:
+                yield from _process_files_batch(
+                    files_to_process, convert_func, self.batch_size
                 )
-                for folder_id in folder_ids
-            ]
-        )
-        for files_batch in file_batches:
-            doc_batch = []
-            for file in files_batch:
-                try:
-                    # Skip files that are shortcuts
-                    if file.get("mimeType") == DRIVE_SHORTCUT_TYPE:
-                        logger.info("Ignoring Drive Shortcut Filetype")
-                        continue
+                files_to_process = []
 
-                    if self.only_org_public:
-                        if "permissions" not in file:
-                            continue
-                        if not any(
-                            permission["type"] == "domain"
-                            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
-
-                    doc_batch.append(
-                        Document(
-                            id=file["webViewLink"],
-                            sections=[
-                                Section(link=file["webViewLink"], text=text_contents)
-                            ],
-                            source=DocumentSource.GOOGLE_DRIVE,
-                            semantic_identifier=file["name"],
-                            doc_updated_at=datetime.fromisoformat(
-                                file["modifiedTime"]
-                            ).astimezone(timezone.utc),
-                            metadata={} if text_contents else {IGNORE_FOR_QA: "True"},
-                            additional_info=file.get("id"),
-                        )
-                    )
-                except Exception as e:
-                    if not self.continue_on_failure:
-                        raise e
-
-                    logger.exception(
-                        "Ran into exception when pulling a file from Google Drive"
-                    )
-
-            yield doc_batch
+        # Process any remaining files
+        if files_to_process:
+            yield from _process_files_batch(
+                files_to_process, convert_func, self.batch_size
+            )
 
     def load_from_state(self) -> GenerateDocumentsOutput:
-        yield from self._fetch_docs_from_drive()
+        try:
+            yield from self._extract_docs_from_google_drive()
+        except Exception as e:
+            if MISSING_SCOPES_ERROR_STR in str(e):
+                raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e
+            raise e
 
     def poll_source(
         self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch
     ) -> GenerateDocumentsOutput:
-        # need to subtract 10 minutes from start time to account for modifiedTime
-        # propogation if a document is modified, it takes some time for the API to
-        # reflect these changes if we do not have an offset, then we may "miss" the
-        # update when polling
-        yield from self._fetch_docs_from_drive(start, end)
+        try:
+            yield from self._extract_docs_from_google_drive(start, end)
+        except Exception as e:
+            if MISSING_SCOPES_ERROR_STR in str(e):
+                raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e
+            raise e
 
+    def _extract_slim_docs_from_google_drive(
+        self,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+    ) -> GenerateSlimDocumentOutput:
+        slim_batch = []
+        for file in self._fetch_drive_items(
+            is_slim=True,
+            start=start,
+            end=end,
+        ):
+            if doc := build_slim_document(file):
+                slim_batch.append(doc)
+            if len(slim_batch) >= SLIM_BATCH_SIZE:
+                yield slim_batch
+                slim_batch = []
+        yield slim_batch
 
-if __name__ == "__main__":
-    import json
-    import os
-
-    service_account_json_path = os.environ.get("GOOGLE_SERVICE_ACCOUNT_KEY_JSON_PATH")
-    if not service_account_json_path:
-        raise ValueError(
-            "Please set GOOGLE_SERVICE_ACCOUNT_KEY_JSON_PATH environment variable"
-        )
-    with open(service_account_json_path) as f:
-        creds = json.load(f)
-
-    credentials_dict = {
-        DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: json.dumps(creds),
-    }
-    delegated_user = os.environ.get("GOOGLE_DRIVE_DELEGATED_USER")
-    if delegated_user:
-        credentials_dict[DB_CREDENTIALS_DICT_DELEGATED_USER_KEY] = delegated_user
-
-    connector = GoogleDriveConnector(include_shared=True, follow_shortcuts=True)
-    connector.load_credentials(credentials_dict)
-    document_batch_generator = connector.load_from_state()
-    for document_batch in document_batch_generator:
-        print(document_batch)
-        break
+    def retrieve_all_slim_documents(
+        self,
+        start: SecondsSinceUnixEpoch | None = None,
+        end: SecondsSinceUnixEpoch | None = None,
+    ) -> GenerateSlimDocumentOutput:
+        try:
+            yield from self._extract_slim_docs_from_google_drive(start, end)
+        except Exception as e:
+            if MISSING_SCOPES_ERROR_STR in str(e):
+                raise PermissionError(ONYX_SCOPE_INSTRUCTIONS) from e
+            raise e

backend/danswer/connectors/google_drive/connector_auth.py

@@ -1,229 +0,0 @@
-import json
-from typing import cast
-from urllib.parse import parse_qs
-from urllib.parse import ParseResult
-from urllib.parse import urlparse
-
-from google.auth.transport.requests import Request  # type: ignore
-from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
-from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
-from google_auth_oauthlib.flow import InstalledAppFlow  # type: ignore
-from sqlalchemy.orm import Session
-
-from danswer.configs.app_configs import ENTERPRISE_EDITION_ENABLED
-from danswer.configs.app_configs import WEB_DOMAIN
-from danswer.configs.constants import DocumentSource
-from danswer.configs.constants import KV_CRED_KEY
-from danswer.configs.constants import KV_GOOGLE_DRIVE_CRED_KEY
-from danswer.configs.constants import KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY
-from danswer.connectors.google_drive.constants import BASE_SCOPES
-from danswer.connectors.google_drive.constants import (
-    DB_CREDENTIALS_DICT_DELEGATED_USER_KEY,
-)
-from danswer.connectors.google_drive.constants import (
-    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
-)
-from danswer.connectors.google_drive.constants import DB_CREDENTIALS_DICT_TOKEN_KEY
-from danswer.connectors.google_drive.constants import FETCH_GROUPS_SCOPES
-from danswer.connectors.google_drive.constants import FETCH_PERMISSIONS_SCOPES
-from danswer.db.credentials import update_credential_json
-from danswer.db.models import User
-from danswer.key_value_store.factory import get_kv_store
-from danswer.server.documents.models import CredentialBase
-from danswer.server.documents.models import GoogleAppCredentials
-from danswer.server.documents.models import GoogleServiceAccountKey
-from danswer.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def build_gdrive_scopes() -> list[str]:
-    base_scopes: list[str] = BASE_SCOPES
-    permissions_scopes: list[str] = FETCH_PERMISSIONS_SCOPES
-    groups_scopes: list[str] = FETCH_GROUPS_SCOPES
-
-    if ENTERPRISE_EDITION_ENABLED:
-        return base_scopes + permissions_scopes + groups_scopes
-    return base_scopes + permissions_scopes
-
-
-def _build_frontend_google_drive_redirect() -> str:
-    return f"{WEB_DOMAIN}/admin/connectors/google-drive/auth/callback"
-
-
-def get_google_drive_creds_for_authorized_user(
-    token_json_str: str, scopes: list[str] = build_gdrive_scopes()
-) -> OAuthCredentials | None:
-    creds_json = json.loads(token_json_str)
-    creds = OAuthCredentials.from_authorized_user_info(creds_json, scopes)
-    if creds.valid:
-        return creds
-
-    if creds.expired and creds.refresh_token:
-        try:
-            creds.refresh(Request())
-            if creds.valid:
-                logger.notice("Refreshed Google Drive tokens.")
-                return creds
-        except Exception as e:
-            logger.exception(f"Failed to refresh google drive access token due to: {e}")
-            return None
-
-    return None
-
-
-def _get_google_drive_creds_for_service_account(
-    service_account_key_json_str: str, scopes: list[str] = build_gdrive_scopes()
-) -> ServiceAccountCredentials | None:
-    service_account_key = json.loads(service_account_key_json_str)
-    creds = ServiceAccountCredentials.from_service_account_info(
-        service_account_key, scopes=scopes
-    )
-    if not creds.valid or not creds.expired:
-        creds.refresh(Request())
-    return creds if creds.valid else None
-
-
-def get_google_drive_creds(
-    credentials: dict[str, str], scopes: list[str] = build_gdrive_scopes()
-) -> tuple[ServiceAccountCredentials | OAuthCredentials, dict[str, str] | None]:
-    oauth_creds = None
-    service_creds = None
-    new_creds_dict = None
-    if DB_CREDENTIALS_DICT_TOKEN_KEY in credentials:
-        access_token_json_str = cast(str, credentials[DB_CREDENTIALS_DICT_TOKEN_KEY])
-        oauth_creds = get_google_drive_creds_for_authorized_user(
-            token_json_str=access_token_json_str, scopes=scopes
-        )
-
-        # tell caller to update token stored in DB if it has changed
-        # (e.g. the token has been refreshed)
-        new_creds_json_str = oauth_creds.to_json() if oauth_creds else ""
-        if new_creds_json_str != access_token_json_str:
-            new_creds_dict = {DB_CREDENTIALS_DICT_TOKEN_KEY: new_creds_json_str}
-
-    elif DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY in credentials:
-        service_account_key_json_str = credentials[
-            DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY
-        ]
-        service_creds = _get_google_drive_creds_for_service_account(
-            service_account_key_json_str=service_account_key_json_str,
-            scopes=scopes,
-        )
-
-        # "Impersonate" a user if one is specified
-        delegated_user_email = cast(
-            str | None, credentials.get(DB_CREDENTIALS_DICT_DELEGATED_USER_KEY)
-        )
-        if delegated_user_email:
-            service_creds = (
-                service_creds.with_subject(delegated_user_email)
-                if service_creds
-                else None
-            )
-
-    creds: ServiceAccountCredentials | OAuthCredentials | None = (
-        oauth_creds or service_creds
-    )
-    if creds is None:
-        raise PermissionError(
-            "Unable to access Google Drive - unknown credential structure."
-        )
-
-    return creds, new_creds_dict
-
-
-def verify_csrf(credential_id: int, state: str) -> None:
-    csrf = get_kv_store().load(KV_CRED_KEY.format(str(credential_id)))
-    if csrf != state:
-        raise PermissionError(
-            "State from Google Drive Connector callback does not match expected"
-        )
-
-
-def get_auth_url(credential_id: int) -> str:
-    creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))
-    credential_json = json.loads(creds_str)
-    flow = InstalledAppFlow.from_client_config(
-        credential_json,
-        scopes=build_gdrive_scopes(),
-        redirect_uri=_build_frontend_google_drive_redirect(),
-    )
-    auth_url, _ = flow.authorization_url(prompt="consent")
-
-    parsed_url = cast(ParseResult, urlparse(auth_url))
-    params = parse_qs(parsed_url.query)
-
-    get_kv_store().store(
-        KV_CRED_KEY.format(credential_id), params.get("state", [None])[0], encrypt=True
-    )  # type: ignore
-    return str(auth_url)
-
-
-def update_credential_access_tokens(
-    auth_code: str,
-    credential_id: int,
-    user: User,
-    db_session: Session,
-) -> OAuthCredentials | None:
-    app_credentials = get_google_app_cred()
-    flow = InstalledAppFlow.from_client_config(
-        app_credentials.model_dump(),
-        scopes=build_gdrive_scopes(),
-        redirect_uri=_build_frontend_google_drive_redirect(),
-    )
-    flow.fetch_token(code=auth_code)
-    creds = flow.credentials
-    token_json_str = creds.to_json()
-    new_creds_dict = {DB_CREDENTIALS_DICT_TOKEN_KEY: token_json_str}
-
-    if not update_credential_json(credential_id, new_creds_dict, user, db_session):
-        return None
-    return creds
-
-
-def build_service_account_creds(
-    source: DocumentSource,
-    delegated_user_email: str | None = None,
-) -> CredentialBase:
-    service_account_key = get_service_account_key()
-
-    credential_dict = {
-        DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY: service_account_key.json(),
-    }
-    if delegated_user_email:
-        credential_dict[DB_CREDENTIALS_DICT_DELEGATED_USER_KEY] = delegated_user_email
-
-    return CredentialBase(
-        credential_json=credential_dict,
-        admin_public=True,
-        source=DocumentSource.GOOGLE_DRIVE,
-    )
-
-
-def get_google_app_cred() -> GoogleAppCredentials:
-    creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_CRED_KEY))
-    return GoogleAppCredentials(**json.loads(creds_str))
-
-
-def upsert_google_app_cred(app_credentials: GoogleAppCredentials) -> None:
-    get_kv_store().store(KV_GOOGLE_DRIVE_CRED_KEY, app_credentials.json(), encrypt=True)
-
-
-def delete_google_app_cred() -> None:
-    get_kv_store().delete(KV_GOOGLE_DRIVE_CRED_KEY)
-
-
-def get_service_account_key() -> GoogleServiceAccountKey:
-    creds_str = str(get_kv_store().load(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY))
-    return GoogleServiceAccountKey(**json.loads(creds_str))
-
-
-def upsert_service_account_key(service_account_key: GoogleServiceAccountKey) -> None:
-    get_kv_store().store(
-        KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY, service_account_key.json(), encrypt=True
-    )
-
-
-def delete_service_account_key() -> None:
-    get_kv_store().delete(KV_GOOGLE_DRIVE_SERVICE_ACCOUNT_KEY)

backend/danswer/connectors/google_drive/constants.py

@@ -1,7 +1,4 @@
-DB_CREDENTIALS_DICT_TOKEN_KEY = "google_drive_tokens"
-DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY = "google_drive_service_account_key"
-DB_CREDENTIALS_DICT_DELEGATED_USER_KEY = "google_drive_delegated_user"
-
-BASE_SCOPES = ["https://www.googleapis.com/auth/drive.readonly"]
-FETCH_PERMISSIONS_SCOPES = ["https://www.googleapis.com/auth/drive.metadata.readonly"]
-FETCH_GROUPS_SCOPES = ["https://www.googleapis.com/auth/cloud-identity.groups.readonly"]
+UNSUPPORTED_FILE_TYPE_CONTENT = ""  # keep empty for now
+DRIVE_FOLDER_TYPE = "application/vnd.google-apps.folder"
+DRIVE_SHORTCUT_TYPE = "application/vnd.google-apps.shortcut"
+DRIVE_FILE_TYPE = "application/vnd.google-apps.file"

backend/danswer/connectors/google_drive/doc_conversion.py

@@ -0,0 +1,197 @@
+import io
+from datetime import datetime
+from datetime import timezone
+
+from googleapiclient.errors import HttpError  # type: ignore
+
+from danswer.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE
+from danswer.configs.constants import DocumentSource
+from danswer.configs.constants import IGNORE_FOR_QA
+from danswer.connectors.google_drive.constants import DRIVE_FOLDER_TYPE
+from danswer.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE
+from danswer.connectors.google_drive.constants import UNSUPPORTED_FILE_TYPE_CONTENT
+from danswer.connectors.google_drive.models import GDriveMimeType
+from danswer.connectors.google_drive.models import GoogleDriveFileType
+from danswer.connectors.google_drive.section_extraction import get_document_sections
+from danswer.connectors.google_utils.resources import GoogleDocsService
+from danswer.connectors.google_utils.resources import GoogleDriveService
+from danswer.connectors.models import Document
+from danswer.connectors.models import Section
+from danswer.connectors.models import SlimDocument
+from danswer.file_processing.extract_file_text import docx_to_text
+from danswer.file_processing.extract_file_text import pptx_to_text
+from danswer.file_processing.extract_file_text import read_pdf_file
+from danswer.file_processing.unstructured import get_unstructured_api_key
+from danswer.file_processing.unstructured import unstructured_to_text
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+# 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",
+]
+
+
+def _extract_sections_basic(
+    file: dict[str, str], service: GoogleDriveService
+) -> list[Section]:
+    mime_type = file["mimeType"]
+    link = file["webViewLink"]
+
+    if mime_type not in set(item.value for item in GDriveMimeType):
+        # Unsupported file types can still have a title, finding this way is still useful
+        return [Section(link=link, text=UNSUPPORTED_FILE_TYPE_CONTENT)]
+
+    try:
+        if mime_type in [
+            GDriveMimeType.DOC.value,
+            GDriveMimeType.PPT.value,
+            GDriveMimeType.SPREADSHEET.value,
+        ]:
+            export_mime_type = (
+                "text/plain"
+                if mime_type != GDriveMimeType.SPREADSHEET.value
+                else "text/csv"
+            )
+            text = (
+                service.files()
+                .export(fileId=file["id"], mimeType=export_mime_type)
+                .execute()
+                .decode("utf-8")
+            )
+            return [Section(link=link, text=text)]
+        elif mime_type in [
+            GDriveMimeType.PLAIN_TEXT.value,
+            GDriveMimeType.MARKDOWN.value,
+        ]:
+            return [
+                Section(
+                    link=link,
+                    text=service.files()
+                    .get_media(fileId=file["id"])
+                    .execute()
+                    .decode("utf-8"),
+                )
+            ]
+        if mime_type in [
+            GDriveMimeType.WORD_DOC.value,
+            GDriveMimeType.POWERPOINT.value,
+            GDriveMimeType.PDF.value,
+        ]:
+            response = service.files().get_media(fileId=file["id"]).execute()
+            if get_unstructured_api_key():
+                return [
+                    Section(
+                        link=link,
+                        text=unstructured_to_text(
+                            file=io.BytesIO(response),
+                            file_name=file.get("name", file["id"]),
+                        ),
+                    )
+                ]
+
+            if mime_type == GDriveMimeType.WORD_DOC.value:
+                return [
+                    Section(link=link, text=docx_to_text(file=io.BytesIO(response)))
+                ]
+            elif mime_type == GDriveMimeType.PDF.value:
+                text, _ = read_pdf_file(file=io.BytesIO(response))
+                return [Section(link=link, text=text)]
+            elif mime_type == GDriveMimeType.POWERPOINT.value:
+                return [
+                    Section(link=link, text=pptx_to_text(file=io.BytesIO(response)))
+                ]
+
+        return [Section(link=link, text=UNSUPPORTED_FILE_TYPE_CONTENT)]
+
+    except Exception:
+        return [Section(link=link, text=UNSUPPORTED_FILE_TYPE_CONTENT)]
+
+
+def convert_drive_item_to_document(
+    file: GoogleDriveFileType,
+    drive_service: GoogleDriveService,
+    docs_service: GoogleDocsService,
+) -> Document | None:
+    try:
+        # Skip files that are shortcuts
+        if file.get("mimeType") == DRIVE_SHORTCUT_TYPE:
+            logger.info("Ignoring Drive Shortcut Filetype")
+            return None
+        # Skip files that are folders
+        if file.get("mimeType") == DRIVE_FOLDER_TYPE:
+            logger.info("Ignoring Drive Folder Filetype")
+            return None
+
+        sections: list[Section] = []
+
+        # Special handling for Google Docs to preserve structure, link
+        # to headers
+        if file.get("mimeType") == GDriveMimeType.DOC.value:
+            try:
+                sections = get_document_sections(docs_service, file["id"])
+            except Exception as e:
+                logger.warning(
+                    f"Ran into exception '{e}' when pulling sections from Google Doc '{file['name']}'."
+                    " Falling back to basic extraction."
+                )
+        # NOTE: this will run for either (1) the above failed or (2) the file is not a Google Doc
+        if not sections:
+            try:
+                # For all other file types just extract the text
+                sections = _extract_sections_basic(file, drive_service)
+
+            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
+                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..."
+                    )
+                    return None
+
+                raise
+        if not sections:
+            return None
+
+        return Document(
+            id=file["webViewLink"],
+            sections=sections,
+            source=DocumentSource.GOOGLE_DRIVE,
+            semantic_identifier=file["name"],
+            doc_updated_at=datetime.fromisoformat(file["modifiedTime"]).astimezone(
+                timezone.utc
+            ),
+            metadata={}
+            if any(section.text for section in sections)
+            else {IGNORE_FOR_QA: "True"},
+            additional_info=file.get("id"),
+        )
+    except Exception as e:
+        if not CONTINUE_ON_CONNECTOR_FAILURE:
+            raise e
+
+        logger.exception("Ran into exception when pulling a file from Google Drive")
+    return None
+
+
+def build_slim_document(file: GoogleDriveFileType) -> SlimDocument | None:
+    # Skip files that are folders or shortcuts
+    if file.get("mimeType") in [DRIVE_FOLDER_TYPE, DRIVE_SHORTCUT_TYPE]:
+        return None
+
+    return SlimDocument(
+        id=file["webViewLink"],
+        perm_sync_data={
+            "doc_id": file.get("id"),
+            "permissions": file.get("permissions", []),
+            "permission_ids": file.get("permissionIds", []),
+            "name": file.get("name"),
+            "owner_email": file.get("owners", [{}])[0].get("emailAddress"),
+        },
+    )

backend/danswer/connectors/google_drive/file_retrieval.py

@@ -0,0 +1,222 @@
+from collections.abc import Callable
+from collections.abc import Iterator
+from datetime import datetime
+from typing import Any
+
+from googleapiclient.discovery import Resource  # type: ignore
+
+from danswer.connectors.google_drive.constants import DRIVE_FOLDER_TYPE
+from danswer.connectors.google_drive.constants import DRIVE_SHORTCUT_TYPE
+from danswer.connectors.google_drive.models import GoogleDriveFileType
+from danswer.connectors.google_utils.google_utils import execute_paginated_retrieval
+from danswer.connectors.interfaces import SecondsSinceUnixEpoch
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+FILE_FIELDS = (
+    "nextPageToken, files(mimeType, id, name, permissions, modifiedTime, webViewLink, "
+    "shortcutDetails, owners(emailAddress))"
+)
+SLIM_FILE_FIELDS = (
+    "nextPageToken, files(mimeType, id, name, permissions(emailAddress, type), "
+    "permissionIds, webViewLink, owners(emailAddress))"
+)
+FOLDER_FIELDS = "nextPageToken, files(id, name, permissions, modifiedTime, webViewLink, shortcutDetails)"
+
+
+def _generate_time_range_filter(
+    start: SecondsSinceUnixEpoch | None = None,
+    end: SecondsSinceUnixEpoch | None = None,
+) -> str:
+    time_range_filter = ""
+    if start is not None:
+        time_start = datetime.utcfromtimestamp(start).isoformat() + "Z"
+        time_range_filter += f" and modifiedTime >= '{time_start}'"
+    if end is not None:
+        time_stop = datetime.utcfromtimestamp(end).isoformat() + "Z"
+        time_range_filter += f" and modifiedTime <= '{time_stop}'"
+    return time_range_filter
+
+
+def _get_folders_in_parent(
+    service: Resource,
+    parent_id: str | None = None,
+) -> Iterator[GoogleDriveFileType]:
+    # Follow shortcuts to folders
+    query = f"(mimeType = '{DRIVE_FOLDER_TYPE}' or mimeType = '{DRIVE_SHORTCUT_TYPE}')"
+    query += " and trashed = false"
+
+    if parent_id:
+        query += f" and '{parent_id}' in parents"
+
+    for file in execute_paginated_retrieval(
+        retrieval_function=service.files().list,
+        list_key="files",
+        continue_on_404_or_403=True,
+        corpora="allDrives",
+        supportsAllDrives=True,
+        includeItemsFromAllDrives=True,
+        fields=FOLDER_FIELDS,
+        q=query,
+    ):
+        yield file
+
+
+def _get_files_in_parent(
+    service: Resource,
+    parent_id: str,
+    start: SecondsSinceUnixEpoch | None = None,
+    end: SecondsSinceUnixEpoch | None = None,
+    is_slim: bool = False,
+) -> Iterator[GoogleDriveFileType]:
+    query = f"mimeType != '{DRIVE_FOLDER_TYPE}' and '{parent_id}' in parents"
+    query += " and trashed = false"
+    query += _generate_time_range_filter(start, end)
+
+    for file in execute_paginated_retrieval(
+        retrieval_function=service.files().list,
+        list_key="files",
+        continue_on_404_or_403=True,
+        corpora="allDrives",
+        supportsAllDrives=True,
+        includeItemsFromAllDrives=True,
+        fields=SLIM_FILE_FIELDS if is_slim else FILE_FIELDS,
+        q=query,
+    ):
+        yield file
+
+
+def crawl_folders_for_files(
+    service: Resource,
+    parent_id: str,
+    traversed_parent_ids: set[str],
+    update_traversed_ids_func: Callable[[str], None],
+    start: SecondsSinceUnixEpoch | None = None,
+    end: SecondsSinceUnixEpoch | None = None,
+) -> Iterator[GoogleDriveFileType]:
+    """
+    This function starts crawling from any folder. It is slower though.
+    """
+    if parent_id in traversed_parent_ids:
+        logger.info(f"Skipping subfolder since already traversed: {parent_id}")
+        return
+
+    found_files = False
+    for file in _get_files_in_parent(
+        service=service,
+        start=start,
+        end=end,
+        parent_id=parent_id,
+    ):
+        found_files = True
+        yield file
+
+    if found_files:
+        update_traversed_ids_func(parent_id)
+
+    for subfolder in _get_folders_in_parent(
+        service=service,
+        parent_id=parent_id,
+    ):
+        logger.info("Fetching all files in subfolder: " + subfolder["name"])
+        yield from crawl_folders_for_files(
+            service=service,
+            parent_id=subfolder["id"],
+            traversed_parent_ids=traversed_parent_ids,
+            update_traversed_ids_func=update_traversed_ids_func,
+            start=start,
+            end=end,
+        )
+
+
+def get_files_in_shared_drive(
+    service: Resource,
+    drive_id: str,
+    is_slim: bool = False,
+    update_traversed_ids_func: Callable[[str], None] = lambda _: None,
+    start: SecondsSinceUnixEpoch | None = None,
+    end: SecondsSinceUnixEpoch | None = None,
+) -> Iterator[GoogleDriveFileType]:
+    # If we know we are going to folder crawl later, we can cache the folders here
+    # Get all folders being queried and add them to the traversed set
+    query = f"mimeType = '{DRIVE_FOLDER_TYPE}'"
+    query += " and trashed = false"
+    found_folders = False
+    for file in execute_paginated_retrieval(
+        retrieval_function=service.files().list,
+        list_key="files",
+        continue_on_404_or_403=True,
+        corpora="drive",
+        driveId=drive_id,
+        supportsAllDrives=True,
+        includeItemsFromAllDrives=True,
+        fields="nextPageToken, files(id)",
+        q=query,
+    ):
+        update_traversed_ids_func(file["id"])
+        found_folders = True
+    if found_folders:
+        update_traversed_ids_func(drive_id)
+
+    # Get all files in the shared drive
+    query = f"mimeType != '{DRIVE_FOLDER_TYPE}'"
+    query += " and trashed = false"
+    query += _generate_time_range_filter(start, end)
+    yield from execute_paginated_retrieval(
+        retrieval_function=service.files().list,
+        list_key="files",
+        continue_on_404_or_403=True,
+        corpora="drive",
+        driveId=drive_id,
+        supportsAllDrives=True,
+        includeItemsFromAllDrives=True,
+        fields=SLIM_FILE_FIELDS if is_slim else FILE_FIELDS,
+        q=query,
+    )
+
+
+def get_all_files_in_my_drive(
+    service: Any,
+    update_traversed_ids_func: Callable,
+    is_slim: bool = False,
+    start: SecondsSinceUnixEpoch | None = None,
+    end: SecondsSinceUnixEpoch | None = None,
+) -> Iterator[GoogleDriveFileType]:
+    # If we know we are going to folder crawl later, we can cache the folders here
+    # Get all folders being queried and add them to the traversed set
+    query = "trashed = false and 'me' in owners"
+    found_folders = False
+    for file in execute_paginated_retrieval(
+        retrieval_function=service.files().list,
+        list_key="files",
+        corpora="user",
+        fields=SLIM_FILE_FIELDS if is_slim else FILE_FIELDS,
+        q=query,
+    ):
+        update_traversed_ids_func(file["id"])
+        found_folders = True
+    if found_folders:
+        update_traversed_ids_func(get_root_folder_id(service))
+
+    # Then get the files
+    query = "trashed = false and 'me' in owners"
+    query += _generate_time_range_filter(start, end)
+    fields = "files(id, name, mimeType, webViewLink, modifiedTime, createdTime)"
+    if not is_slim:
+        fields += ", files(permissions, permissionIds, owners)"
+
+    yield from execute_paginated_retrieval(
+        retrieval_function=service.files().list,
+        list_key="files",
+        corpora="user",
+        fields=SLIM_FILE_FIELDS if is_slim else FILE_FIELDS,
+        q=query,
+    )
+
+
+# Just in case we need to get the root folder id
+def get_root_folder_id(service: Resource) -> str:
+    # we dont paginate here because there is only one root folder per user
+    # https://developers.google.com/drive/api/guides/v2-to-v3-reference
+    return service.files().get(fileId="root", fields="id").execute()["id"]

backend/danswer/connectors/google_drive/models.py

@@ -0,0 +1,18 @@
+from enum import Enum
+from typing import Any
+
+
+class GDriveMimeType(str, Enum):
+    DOC = "application/vnd.google-apps.document"
+    SPREADSHEET = "application/vnd.google-apps.spreadsheet"
+    PDF = "application/pdf"
+    WORD_DOC = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    PPT = "application/vnd.google-apps.presentation"
+    POWERPOINT = (
+        "application/vnd.openxmlformats-officedocument.presentationml.presentation"
+    )
+    PLAIN_TEXT = "text/plain"
+    MARKDOWN = "text/markdown"
+
+
+GoogleDriveFileType = dict[str, Any]

backend/danswer/connectors/google_drive/section_extraction.py

@@ -0,0 +1,105 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from danswer.connectors.google_utils.resources import GoogleDocsService
+from danswer.connectors.models import Section
+
+
+class CurrentHeading(BaseModel):
+    id: str
+    text: str
+
+
+def _build_gdoc_section_link(doc_id: str, heading_id: str) -> str:
+    """Builds a Google Doc link that jumps to a specific heading"""
+    # NOTE: doesn't support docs with multiple tabs atm, if we need that ask
+    # @Chris
+    return (
+        f"https://docs.google.com/document/d/{doc_id}/edit?tab=t.0#heading={heading_id}"
+    )
+
+
+def _extract_id_from_heading(paragraph: dict[str, Any]) -> str:
+    """Extracts the id from a heading paragraph element"""
+    return paragraph["paragraphStyle"]["headingId"]
+
+
+def _extract_text_from_paragraph(paragraph: dict[str, Any]) -> str:
+    """Extracts the text content from a paragraph element"""
+    text_elements = []
+    for element in paragraph.get("elements", []):
+        if "textRun" in element:
+            text_elements.append(element["textRun"].get("content", ""))
+    return "".join(text_elements)
+
+
+def get_document_sections(
+    docs_service: GoogleDocsService,
+    doc_id: str,
+) -> list[Section]:
+    """Extracts sections from a Google Doc, including their headings and content"""
+    # Fetch the document structure
+    doc = docs_service.documents().get(documentId=doc_id).execute()
+
+    # Get the content
+    content = doc.get("body", {}).get("content", [])
+
+    sections: list[Section] = []
+    current_section: list[str] = []
+    current_heading: CurrentHeading | None = None
+
+    for element in content:
+        if "paragraph" not in element:
+            continue
+
+        paragraph = element["paragraph"]
+
+        # Check if this is a heading
+        if (
+            "paragraphStyle" in paragraph
+            and "namedStyleType" in paragraph["paragraphStyle"]
+        ):
+            style = paragraph["paragraphStyle"]["namedStyleType"]
+            is_heading = style.startswith("HEADING_")
+            is_title = style.startswith("TITLE")
+
+            if is_heading or is_title:
+                # If we were building a previous section, add it to sections list
+                if current_heading is not None and current_section:
+                    heading_text = current_heading.text
+                    section_text = f"{heading_text}\n" + "\n".join(current_section)
+                    sections.append(
+                        Section(
+                            text=section_text.strip(),
+                            link=_build_gdoc_section_link(doc_id, current_heading.id),
+                        )
+                    )
+                    current_section = []
+
+                # Start new heading
+                heading_id = _extract_id_from_heading(paragraph)
+                heading_text = _extract_text_from_paragraph(paragraph)
+                current_heading = CurrentHeading(
+                    id=heading_id,
+                    text=heading_text,
+                )
+                continue
+
+        # Add content to current section
+        if current_heading is not None:
+            text = _extract_text_from_paragraph(paragraph)
+            if text.strip():
+                current_section.append(text)
+
+    # Don't forget to add the last section
+    if current_heading is not None and current_section:
+        section_text = f"{current_heading.text}\n" + "\n".join(current_section)
+        sections.append(
+            Section(
+                text=section_text.strip(),
+                link=_build_gdoc_section_link(doc_id, current_heading.id),
+            )
+        )
+
+    return sections

backend/danswer/connectors/google_utils/google_auth.py

@@ -0,0 +1,107 @@
+import json
+from typing import cast
+
+from google.auth.transport.requests import Request  # type: ignore
+from google.oauth2.credentials import Credentials as OAuthCredentials  # type: ignore
+from google.oauth2.service_account import Credentials as ServiceAccountCredentials  # type: ignore
+
+from danswer.configs.constants import DocumentSource
+from danswer.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY,
+)
+from danswer.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_DICT_TOKEN_KEY,
+)
+from danswer.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_PRIMARY_ADMIN_KEY,
+)
+from danswer.connectors.google_utils.shared_constants import (
+    GOOGLE_SCOPES,
+)
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def get_google_oauth_creds(
+    token_json_str: str, source: DocumentSource
+) -> OAuthCredentials | None:
+    creds_json = json.loads(token_json_str)
+    creds = OAuthCredentials.from_authorized_user_info(
+        info=creds_json,
+        scopes=GOOGLE_SCOPES[source],
+    )
+    if creds.valid:
+        return creds
+
+    if creds.expired and creds.refresh_token:
+        try:
+            creds.refresh(Request())
+            if creds.valid:
+                logger.notice("Refreshed Google Drive tokens.")
+                return creds
+        except Exception:
+            logger.exception("Failed to refresh google drive access token due to:")
+            return None
+
+    return None
+
+
+def get_google_creds(
+    credentials: dict[str, str],
+    source: DocumentSource,
+) -> tuple[ServiceAccountCredentials | OAuthCredentials, dict[str, str] | None]:
+    """Checks for two different types of credentials.
+    (1) A credential which holds a token acquired via a user going thorough
+    the Google OAuth flow.
+    (2) A credential which holds a service account key JSON file, which
+    can then be used to impersonate any user in the workspace.
+    """
+    oauth_creds = None
+    service_creds = None
+    new_creds_dict = None
+    if DB_CREDENTIALS_DICT_TOKEN_KEY in credentials:
+        # OAUTH
+        access_token_json_str = cast(str, credentials[DB_CREDENTIALS_DICT_TOKEN_KEY])
+        oauth_creds = get_google_oauth_creds(
+            token_json_str=access_token_json_str, source=source
+        )
+
+        # tell caller to update token stored in DB if it has changed
+        # (e.g. the token has been refreshed)
+        new_creds_json_str = oauth_creds.to_json() if oauth_creds else ""
+        if new_creds_json_str != access_token_json_str:
+            new_creds_dict = {
+                DB_CREDENTIALS_DICT_TOKEN_KEY: new_creds_json_str,
+                DB_CREDENTIALS_PRIMARY_ADMIN_KEY: credentials[
+                    DB_CREDENTIALS_PRIMARY_ADMIN_KEY
+                ],
+            }
+    elif DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY in credentials:
+        # SERVICE ACCOUNT
+        service_account_key_json_str = credentials[
+            DB_CREDENTIALS_DICT_SERVICE_ACCOUNT_KEY
+        ]
+        service_account_key = json.loads(service_account_key_json_str)
+
+        service_creds = ServiceAccountCredentials.from_service_account_info(
+            service_account_key, scopes=GOOGLE_SCOPES[source]
+        )
+
+        if not service_creds.valid or not service_creds.expired:
+            service_creds.refresh(Request())
+
+        if not service_creds.valid:
+            raise PermissionError(
+                f"Unable to access {source} - service account credentials are invalid."
+            )
+
+    creds: ServiceAccountCredentials | OAuthCredentials | None = (
+        oauth_creds or service_creds
+    )
+    if creds is None:
+        raise PermissionError(
+            f"Unable to access {source} - unknown credential structure."
+        )
+
+    return creds, new_creds_dict