specifically apply flex none to in progress!

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

@@ -65,7 +65,6 @@ jobs:
             NEXT_PUBLIC_POSTHOG_KEY=${{ secrets.POSTHOG_KEY }}
             NEXT_PUBLIC_POSTHOG_HOST=${{ secrets.POSTHOG_HOST }}
             NEXT_PUBLIC_SENTRY_DSN=${{ secrets.SENTRY_DSN }}
-            NEXT_PUBLIC_GTM_ENABLED=true
           # needed due to weird interactions with the builds for different platforms
           no-cache: true
           labels: ${{ steps.meta.outputs.labels }}

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

@@ -13,10 +13,7 @@ on:
 env:
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
   SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
-  CONFLUENCE_TEST_SPACE_URL: ${{ secrets.CONFLUENCE_TEST_SPACE_URL }}
-  CONFLUENCE_USER_NAME: ${{ secrets.CONFLUENCE_USER_NAME }}
-  CONFLUENCE_ACCESS_TOKEN: ${{ secrets.CONFLUENCE_ACCESS_TOKEN }}
-  
+
 jobs:
   integration-tests:
     # See https://runs-on.com/runners/linux/
@@ -198,13 +195,9 @@ jobs:
             -e API_SERVER_HOST=api_server \
             -e OPENAI_API_KEY=${OPENAI_API_KEY} \
             -e SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN} \
-            -e CONFLUENCE_TEST_SPACE_URL=${CONFLUENCE_TEST_SPACE_URL} \
-            -e CONFLUENCE_USER_NAME=${CONFLUENCE_USER_NAME} \
-            -e CONFLUENCE_ACCESS_TOKEN=${CONFLUENCE_ACCESS_TOKEN} \
             -e TEST_WEB_HOSTNAME=test-runner \
             danswer/danswer-integration:test \
-            /app/tests/integration/tests \
-            /app/tests/integration/connector_job_tests
+            /app/tests/integration/tests
         continue-on-error: true
         id: run_tests
 

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

@@ -1,225 +0,0 @@
-name: Run Chromatic Tests
-concurrency:
-  group: Run-Chromatic-Tests-${{ github.workflow }}-${{ github.head_ref || github.event.workflow_run.head_branch || github.run_id }}
-  cancel-in-progress: true
-
-on: push
-
-env:
-  OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-  SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
-
-jobs:
-  playwright-tests:
-    name: Playwright Tests
-
-    # See https://runs-on.com/runners/linux/
-    runs-on: [runs-on,runner=8cpu-linux-x64,ram=16,"run-id=${{ github.run_id }}"]
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-          
-      - 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
-      - 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
-        
-      - name: Setup node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 22
-
-      - name: Install node dependencies
-        working-directory: ./web
-        run: npm ci
-
-      - name: Install playwright browsers
-        working-directory: ./web
-        run: npx playwright install --with-deps
-        
-      - 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 }}
-
-      # tag every docker image with "test" so that we can spin up the correct set
-      # of images during testing
-      
-      # we use the runs-on cache for docker builds
-      # in conjunction with runs-on runners, it has better speed and unlimited caching
-      # https://runs-on.com/caching/s3-cache-for-github-actions/
-      # https://runs-on.com/caching/docker/
-      # https://github.com/moby/buildkit#s3-cache-experimental
-      
-      # images are built and run locally for testing purposes. Not pushed.
-
-      - name: Build Web Docker image
-        uses: ./.github/actions/custom-build-and-push
-        with:
-          context: ./web
-          file: ./web/Dockerfile
-          platforms: linux/amd64
-          tags: danswer/danswer-web-server:test
-          push: false
-          load: true
-          cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/web-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }}
-          cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/web-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max
-
-      - name: Build Backend Docker image
-        uses: ./.github/actions/custom-build-and-push
-        with:
-          context: ./backend
-          file: ./backend/Dockerfile
-          platforms: linux/amd64
-          tags: danswer/danswer-backend:test
-          push: false
-          load: true
-          cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }}
-          cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max
-
-      - name: Build Model Server Docker image
-        uses: ./.github/actions/custom-build-and-push
-        with:
-          context: ./backend
-          file: ./backend/Dockerfile.model_server
-          platforms: linux/amd64
-          tags: danswer/danswer-model-server:test
-          push: false
-          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: Start Docker containers
-        run: |
-          cd deployment/docker_compose
-          ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=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
-
-      - name: Wait for service to be ready
-        run: |
-          echo "Starting wait-for-service script..."
-          
-          docker logs -f danswer-stack-api_server-1 &
-
-          start_time=$(date +%s)
-          timeout=300  # 5 minutes in seconds
-          
-          while true; do
-            current_time=$(date +%s)
-            elapsed_time=$((current_time - start_time))
-            
-            if [ $elapsed_time -ge $timeout ]; then
-              echo "Timeout reached. Service did not become ready in 5 minutes."
-              exit 1
-            fi
-            
-            # Use curl with error handling to ignore specific exit code 56
-            response=$(curl -s -o /dev/null -w "%{http_code}" http://localhost:8080/health || echo "curl_error")
-            
-            if [ "$response" = "200" ]; then
-              echo "Service is ready!"
-              break
-            elif [ "$response" = "curl_error" ]; then
-              echo "Curl encountered an error, possibly exit code 56. Continuing to retry..."
-            else
-              echo "Service not ready yet (HTTP status $response). Retrying in 5 seconds..."
-            fi
-            
-            sleep 5
-          done
-          echo "Finished waiting for service."
-
-      - name: Run pytest playwright test init
-        working-directory: ./backend
-        env: 
-          PYTEST_IGNORE_SKIP: true
-        run: pytest -s tests/integration/tests/playwright/test_playwright.py
-
-      - name: Run Playwright tests
-        working-directory: ./web
-        run: npx playwright test
-
-      - uses: actions/upload-artifact@v4
-        if: always()
-        with:
-          # Chromatic automatically defaults to the test-results directory.
-          # Replace with the path to your custom directory and adjust the CHROMATIC_ARCHIVE_LOCATION environment variable accordingly.
-          name: test-results
-          path: ./web/test-results
-          retention-days: 30
-                    
-      # 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: Upload logs
-        if: success() || failure()
-        uses: actions/upload-artifact@v4
-        with:
-          name: docker-logs
-          path: ${{ 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
-
-  chromatic-tests:
-    name: Chromatic Tests
-    
-    needs: playwright-tests
-    runs-on: [runs-on,runner=8cpu-linux-x64,ram=16,"run-id=${{ github.run_id }}"]
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-          
-      - name: Setup node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 22
-          
-      - name: Install node dependencies
-        working-directory: ./web
-        run: npm ci
-        
-      - name: Download Playwright test results
-        uses: actions/download-artifact@v4
-        with:
-          name: test-results
-          path: ./web/test-results
-          
-      - name: Run Chromatic
-        uses: chromaui/action@latest
-        with:
-          playwright: true
-          projectToken: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
-          workingDir: ./web
-        env: 
-          CHROMATIC_ARCHIVE_LOCATION: ./test-results

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

@@ -23,6 +23,21 @@ jobs:
       with:
         version: v3.14.4
       
+    - 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
+    - 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
+
     - name: Set up chart-testing
       uses: helm/chart-testing-action@v2.6.1
 
@@ -37,22 +52,6 @@ jobs:
           echo "changed=true" >> "$GITHUB_OUTPUT"
         fi
 
-#     rkuo: I don't think we need python?
-#     - 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
-#     - 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
-
     # lint all charts if any changes were detected
     - name: Run chart-testing (lint)
       if: steps.list-changed.outputs.changed == 'true'

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

@@ -20,12 +20,9 @@ env:
   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_TEST_USER_1: ${{ secrets.GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR_TEST_USER_1 }}
   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 }}
-  # Slab
-  SLAB_BOT_TOKEN: ${{ secrets.SLAB_BOT_TOKEN }}
 
 jobs:
   connectors-check:

.gitignore

@@ -7,4 +7,3 @@
 .vscode/
 *.sw?
 /backend/tests/regression/answer_quality/search_test_config.yaml
-/web/test-results/

.vscode/launch.template.jsonc

@@ -203,7 +203,7 @@
                 "--loglevel=INFO",
                 "--hostname=light@%n",
                 "-Q",
-                "vespa_metadata_sync,connector_deletion,doc_permissions_upsert",
+                "vespa_metadata_sync,connector_deletion",
             ],
             "presentation": {
 				 "group": "2",
@@ -232,7 +232,7 @@
                 "--loglevel=INFO",
                 "--hostname=heavy@%n",
                 "-Q",
-                "connector_pruning,connector_doc_permissions_sync,connector_external_group_sync",
+                "connector_pruning",
             ],
             "presentation": {
 				 "group": "2",

CONTRIBUTING.md

@@ -32,7 +32,7 @@ To contribute to this project, please follow the
 When opening a pull request, mention related issues and feel free to tag relevant maintainers.
 
 Before creating a pull request please make sure that the new changes conform to the formatting and linting requirements.
-See the [Formatting and Linting](#formatting-and-linting) section for how to run these checks locally.
+See the [Formatting and Linting](#-formatting-and-linting) section for how to run these checks locally.
 
 
 ### Getting Help 🙋

README.md

@@ -12,7 +12,7 @@
 <a href="https://docs.danswer.dev/" target="_blank">
     <img src="https://img.shields.io/badge/docs-view-blue" alt="Documentation">
 </a>
-<a href="https://join.slack.com/t/danswer/shared_invite/zt-2twesxdr6-5iQitKZQpgq~hYIZ~dv3KA" target="_blank">
+<a href="https://join.slack.com/t/danswer/shared_invite/zt-2lcmqw703-071hBuZBfNEOGUsLa5PXvQ" target="_blank">
     <img src="https://img.shields.io/badge/slack-join-blue.svg?logo=slack" alt="Slack">
 </a>
 <a href="https://discord.gg/TDJ59cGV2X" target="_blank">
@@ -135,7 +135,7 @@ Looking to contribute? Please check out the [Contribution Guide](CONTRIBUTING.md
 
 ## ✨Contributors
 
-<a href="https://github.com/danswer-ai/danswer/graphs/contributors">
+<a href="https://github.com/aryn-ai/sycamore/graphs/contributors">
   <img alt="contributors" src="https://contrib.rocks/image?repo=danswer-ai/danswer"/>
 </a>
 

backend/Dockerfile

@@ -73,7 +73,6 @@ RUN apt-get update && \
     rm -rf /var/lib/apt/lists/* && \
     rm -f /usr/local/lib/python3.11/site-packages/tornado/test/test.key
 
-
 # Pre-downloading models for setups with limited egress
 RUN python -c "from tokenizers import Tokenizer; \
 Tokenizer.from_pretrained('nomic-ai/nomic-embed-text-v1')"

backend/alembic/env.py

@@ -1,5 +1,5 @@
 from sqlalchemy.engine.base import Connection
-from typing import Literal
+from typing import Any
 import asyncio
 from logging.config import fileConfig
 import logging
@@ -8,7 +8,6 @@ from alembic import context
 from sqlalchemy import pool
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.sql import text
-from sqlalchemy.sql.schema import SchemaItem
 
 from shared_configs.configs import MULTI_TENANT
 from danswer.db.engine import build_connection_string
@@ -36,18 +35,7 @@ logger = logging.getLogger(__name__)
 
 
 def include_object(
-    object: SchemaItem,
-    name: str | None,
-    type_: Literal[
-        "schema",
-        "table",
-        "column",
-        "index",
-        "unique_constraint",
-        "foreign_key_constraint",
-    ],
-    reflected: bool,
-    compare_to: SchemaItem | None,
+    object: Any, name: str, type_: str, reflected: bool, compare_to: Any
 ) -> bool:
     """
     Determines whether a database object should be included in migrations.

backend/alembic/versions/177de57c21c9_display_custom_llm_models.py

@@ -1,59 +0,0 @@
-"""display custom llm models
-
-Revision ID: 177de57c21c9
-Revises: 4ee1287bd26a
-Create Date: 2024-11-21 11:49:04.488677
-
-"""
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-from sqlalchemy import and_
-
-revision = "177de57c21c9"
-down_revision = "4ee1287bd26a"
-branch_labels = None
-depends_on = None
-depends_on = None
-
-
-def upgrade() -> None:
-    conn = op.get_bind()
-    llm_provider = sa.table(
-        "llm_provider",
-        sa.column("id", sa.Integer),
-        sa.column("provider", sa.String),
-        sa.column("model_names", postgresql.ARRAY(sa.String)),
-        sa.column("display_model_names", postgresql.ARRAY(sa.String)),
-    )
-
-    excluded_providers = ["openai", "bedrock", "anthropic", "azure"]
-
-    providers_to_update = sa.select(
-        llm_provider.c.id,
-        llm_provider.c.model_names,
-        llm_provider.c.display_model_names,
-    ).where(
-        and_(
-            ~llm_provider.c.provider.in_(excluded_providers),
-            llm_provider.c.model_names.isnot(None),
-        )
-    )
-
-    results = conn.execute(providers_to_update).fetchall()
-
-    for provider_id, model_names, display_model_names in results:
-        if display_model_names is None:
-            display_model_names = []
-
-        combined_model_names = list(set(display_model_names + model_names))
-        update_stmt = (
-            llm_provider.update()
-            .where(llm_provider.c.id == provider_id)
-            .values(display_model_names=combined_model_names)
-        )
-        conn.execute(update_stmt)
-
-
-def downgrade() -> None:
-    pass

backend/alembic/versions/26b931506ecb_default_chosen_assistants_to_none.py

@@ -1,68 +0,0 @@
-"""default chosen assistants to none
-
-Revision ID: 26b931506ecb
-Revises: 2daa494a0851
-Create Date: 2024-11-12 13:23:29.858995
-
-"""
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-# revision identifiers, used by Alembic.
-revision = "26b931506ecb"
-down_revision = "2daa494a0851"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "user", sa.Column("chosen_assistants_new", postgresql.JSONB(), nullable=True)
-    )
-
-    op.execute(
-        """
-    UPDATE "user"
-    SET chosen_assistants_new =
-        CASE
-            WHEN chosen_assistants = '[-2, -1, 0]' THEN NULL
-            ELSE chosen_assistants
-        END
-    """
-    )
-
-    op.drop_column("user", "chosen_assistants")
-
-    op.alter_column(
-        "user", "chosen_assistants_new", new_column_name="chosen_assistants"
-    )
-
-
-def downgrade() -> None:
-    op.add_column(
-        "user",
-        sa.Column(
-            "chosen_assistants_old",
-            postgresql.JSONB(),
-            nullable=False,
-            server_default="[-2, -1, 0]",
-        ),
-    )
-
-    op.execute(
-        """
-    UPDATE "user"
-    SET chosen_assistants_old =
-        CASE
-            WHEN chosen_assistants IS NULL THEN '[-2, -1, 0]'::jsonb
-            ELSE chosen_assistants
-        END
-    """
-    )
-
-    op.drop_column("user", "chosen_assistants")
-
-    op.alter_column(
-        "user", "chosen_assistants_old", new_column_name="chosen_assistants"
-    )

backend/alembic/versions/2daa494a0851_add_group_sync_time.py

@@ -1,30 +0,0 @@
-"""add-group-sync-time
-
-Revision ID: 2daa494a0851
-Revises: c0fd6e4da83a
-Create Date: 2024-11-11 10:57:22.991157
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "2daa494a0851"
-down_revision = "c0fd6e4da83a"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "connector_credential_pair",
-        sa.Column(
-            "last_time_external_group_sync",
-            sa.DateTime(timezone=True),
-            nullable=True,
-        ),
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("connector_credential_pair", "last_time_external_group_sync")

backend/alembic/versions/47e5bef3a1d7_add_persona_categories.py

@@ -1,45 +0,0 @@
-"""add persona categories
-
-Revision ID: 47e5bef3a1d7
-Revises: dfbe9e93d3c7
-Create Date: 2024-11-05 18:55:02.221064
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "47e5bef3a1d7"
-down_revision = "dfbe9e93d3c7"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Create the persona_category table
-    op.create_table(
-        "persona_category",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("name", sa.String(), nullable=False),
-        sa.Column("description", sa.String(), nullable=True),
-        sa.PrimaryKeyConstraint("id"),
-        sa.UniqueConstraint("name"),
-    )
-
-    # Add category_id to persona table
-    op.add_column("persona", sa.Column("category_id", sa.Integer(), nullable=True))
-    op.create_foreign_key(
-        "fk_persona_category",
-        "persona",
-        "persona_category",
-        ["category_id"],
-        ["id"],
-        ondelete="SET NULL",
-    )
-
-
-def downgrade() -> None:
-    op.drop_constraint("fk_persona_category", "persona", type_="foreignkey")
-    op.drop_column("persona", "category_id")
-    op.drop_table("persona_category")

backend/alembic/versions/4ee1287bd26a_add_multiple_slack_bot_support.py

@@ -1,280 +0,0 @@
-"""add_multiple_slack_bot_support
-
-Revision ID: 4ee1287bd26a
-Revises: 47e5bef3a1d7
-Create Date: 2024-11-06 13:15:53.302644
-
-"""
-import logging
-from typing import cast
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.orm import Session
-from danswer.key_value_store.factory import get_kv_store
-from danswer.db.models import SlackBot
-from sqlalchemy.dialects import postgresql
-
-# revision identifiers, used by Alembic.
-revision = "4ee1287bd26a"
-down_revision = "47e5bef3a1d7"
-branch_labels: None = None
-depends_on: None = None
-
-# Configure logging
-logger = logging.getLogger("alembic.runtime.migration")
-logger.setLevel(logging.INFO)
-
-
-def upgrade() -> None:
-    logger.info(f"{revision}: create_table: slack_bot")
-    # Create new slack_bot table
-    op.create_table(
-        "slack_bot",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("name", sa.String(), nullable=False),
-        sa.Column("enabled", sa.Boolean(), nullable=False, server_default="true"),
-        sa.Column("bot_token", sa.LargeBinary(), nullable=False),
-        sa.Column("app_token", sa.LargeBinary(), nullable=False),
-        sa.PrimaryKeyConstraint("id"),
-        sa.UniqueConstraint("bot_token"),
-        sa.UniqueConstraint("app_token"),
-    )
-
-    # # Create new slack_channel_config table
-    op.create_table(
-        "slack_channel_config",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("slack_bot_id", sa.Integer(), nullable=True),
-        sa.Column("persona_id", sa.Integer(), nullable=True),
-        sa.Column("channel_config", postgresql.JSONB(), nullable=False),
-        sa.Column("response_type", sa.String(), nullable=False),
-        sa.Column(
-            "enable_auto_filters", sa.Boolean(), nullable=False, server_default="false"
-        ),
-        sa.ForeignKeyConstraint(
-            ["slack_bot_id"],
-            ["slack_bot.id"],
-        ),
-        sa.ForeignKeyConstraint(
-            ["persona_id"],
-            ["persona.id"],
-        ),
-        sa.PrimaryKeyConstraint("id"),
-    )
-
-    # Handle existing Slack bot tokens first
-    logger.info(f"{revision}: Checking for existing Slack bot.")
-    bot_token = None
-    app_token = None
-    first_row_id = None
-
-    try:
-        tokens = cast(dict, get_kv_store().load("slack_bot_tokens_config_key"))
-    except Exception:
-        logger.warning("No existing Slack bot tokens found.")
-        tokens = {}
-
-    bot_token = tokens.get("bot_token")
-    app_token = tokens.get("app_token")
-
-    if bot_token and app_token:
-        logger.info(f"{revision}: Found bot and app tokens.")
-
-        session = Session(bind=op.get_bind())
-        new_slack_bot = SlackBot(
-            name="Slack Bot (Migrated)",
-            enabled=True,
-            bot_token=bot_token,
-            app_token=app_token,
-        )
-        session.add(new_slack_bot)
-        session.commit()
-        first_row_id = new_slack_bot.id
-
-    # Create a default bot if none exists
-    # This is in case there are no slack tokens but there are channels configured
-    op.execute(
-        sa.text(
-            """
-            INSERT INTO slack_bot (name, enabled, bot_token, app_token)
-            SELECT 'Default Bot', true, '', ''
-            WHERE NOT EXISTS (SELECT 1 FROM slack_bot)
-            RETURNING id;
-            """
-        )
-    )
-
-    # Get the bot ID to use (either from existing migration or newly created)
-    bot_id_query = sa.text(
-        """
-        SELECT COALESCE(
-            :first_row_id,
-            (SELECT id FROM slack_bot ORDER BY id ASC LIMIT 1)
-        ) as bot_id;
-        """
-    )
-    result = op.get_bind().execute(bot_id_query, {"first_row_id": first_row_id})
-    bot_id = result.scalar()
-
-    # CTE (Common Table Expression) that transforms the old slack_bot_config table data
-    # This splits up the channel_names into their own rows
-    channel_names_cte = """
-        WITH channel_names AS (
-            SELECT
-                sbc.id as config_id,
-                sbc.persona_id,
-                sbc.response_type,
-                sbc.enable_auto_filters,
-                jsonb_array_elements_text(sbc.channel_config->'channel_names') as channel_name,
-                sbc.channel_config->>'respond_tag_only' as respond_tag_only,
-                sbc.channel_config->>'respond_to_bots' as respond_to_bots,
-                sbc.channel_config->'respond_member_group_list' as respond_member_group_list,
-                sbc.channel_config->'answer_filters' as answer_filters,
-                sbc.channel_config->'follow_up_tags' as follow_up_tags
-            FROM slack_bot_config sbc
-        )
-    """
-
-    # Insert the channel names into the new slack_channel_config table
-    insert_statement = """
-        INSERT INTO slack_channel_config (
-            slack_bot_id,
-            persona_id,
-            channel_config,
-            response_type,
-            enable_auto_filters
-        )
-        SELECT
-            :bot_id,
-            channel_name.persona_id,
-            jsonb_build_object(
-                'channel_name', channel_name.channel_name,
-                'respond_tag_only',
-                COALESCE((channel_name.respond_tag_only)::boolean, false),
-                'respond_to_bots',
-                COALESCE((channel_name.respond_to_bots)::boolean, false),
-                'respond_member_group_list',
-                COALESCE(channel_name.respond_member_group_list, '[]'::jsonb),
-                'answer_filters',
-                COALESCE(channel_name.answer_filters, '[]'::jsonb),
-                'follow_up_tags',
-                COALESCE(channel_name.follow_up_tags, '[]'::jsonb)
-            ),
-            channel_name.response_type,
-            channel_name.enable_auto_filters
-        FROM channel_names channel_name;
-    """
-
-    op.execute(sa.text(channel_names_cte + insert_statement).bindparams(bot_id=bot_id))
-
-    # Clean up old tokens if they existed
-    try:
-        if bot_token and app_token:
-            logger.info(f"{revision}: Removing old bot and app tokens.")
-            get_kv_store().delete("slack_bot_tokens_config_key")
-    except Exception:
-        logger.warning("tried to delete tokens in dynamic config but failed")
-    # Rename the table
-    op.rename_table(
-        "slack_bot_config__standard_answer_category",
-        "slack_channel_config__standard_answer_category",
-    )
-
-    # Rename the column
-    op.alter_column(
-        "slack_channel_config__standard_answer_category",
-        "slack_bot_config_id",
-        new_column_name="slack_channel_config_id",
-    )
-
-    # Drop the table with CASCADE to handle dependent objects
-    op.execute("DROP TABLE slack_bot_config CASCADE")
-
-    logger.info(f"{revision}: Migration complete.")
-
-
-def downgrade() -> None:
-    # Recreate the old slack_bot_config table
-    op.create_table(
-        "slack_bot_config",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("persona_id", sa.Integer(), nullable=True),
-        sa.Column("channel_config", postgresql.JSONB(), nullable=False),
-        sa.Column("response_type", sa.String(), nullable=False),
-        sa.Column("enable_auto_filters", sa.Boolean(), nullable=False),
-        sa.ForeignKeyConstraint(
-            ["persona_id"],
-            ["persona.id"],
-        ),
-        sa.PrimaryKeyConstraint("id"),
-    )
-
-    # Migrate data back to the old format
-    # Group by persona_id to combine channel names back into arrays
-    op.execute(
-        sa.text(
-            """
-            INSERT INTO slack_bot_config (
-                persona_id,
-                channel_config,
-                response_type,
-                enable_auto_filters
-            )
-            SELECT DISTINCT ON (persona_id)
-                persona_id,
-                jsonb_build_object(
-                    'channel_names', (
-                        SELECT jsonb_agg(c.channel_config->>'channel_name')
-                        FROM slack_channel_config c
-                        WHERE c.persona_id = scc.persona_id
-                    ),
-                    'respond_tag_only', (channel_config->>'respond_tag_only')::boolean,
-                    'respond_to_bots', (channel_config->>'respond_to_bots')::boolean,
-                    'respond_member_group_list', channel_config->'respond_member_group_list',
-                    'answer_filters', channel_config->'answer_filters',
-                    'follow_up_tags', channel_config->'follow_up_tags'
-                ),
-                response_type,
-                enable_auto_filters
-            FROM slack_channel_config scc
-            WHERE persona_id IS NOT NULL;
-            """
-        )
-    )
-
-    # Rename the table back
-    op.rename_table(
-        "slack_channel_config__standard_answer_category",
-        "slack_bot_config__standard_answer_category",
-    )
-
-    # Rename the column back
-    op.alter_column(
-        "slack_bot_config__standard_answer_category",
-        "slack_channel_config_id",
-        new_column_name="slack_bot_config_id",
-    )
-
-    # Try to save the first bot's tokens back to KV store
-    try:
-        first_bot = (
-            op.get_bind()
-            .execute(
-                sa.text(
-                    "SELECT bot_token, app_token FROM slack_bot ORDER BY id LIMIT 1"
-                )
-            )
-            .first()
-        )
-        if first_bot and first_bot.bot_token and first_bot.app_token:
-            tokens = {
-                "bot_token": first_bot.bot_token,
-                "app_token": first_bot.app_token,
-            }
-            get_kv_store().store("slack_bot_tokens_config_key", tokens)
-    except Exception:
-        logger.warning("Failed to save tokens back to KV store")
-
-    # Drop the new tables in reverse order
-    op.drop_table("slack_channel_config")
-    op.drop_table("slack_bot")

backend/alembic/versions/6d562f86c78b_remove_default_bot.py

@@ -1,45 +0,0 @@
-"""remove default bot
-
-Revision ID: 6d562f86c78b
-Revises: 177de57c21c9
-Create Date: 2024-11-22 11:51:29.331336
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "6d562f86c78b"
-down_revision = "177de57c21c9"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.execute(
-        sa.text(
-            """
-            DELETE FROM slack_bot
-            WHERE name = 'Default Bot'
-            AND bot_token = ''
-            AND app_token = ''
-            AND NOT EXISTS (
-                SELECT 1 FROM slack_channel_config
-                WHERE slack_channel_config.slack_bot_id = slack_bot.id
-            )
-            """
-        )
-    )
-
-
-def downgrade() -> None:
-    op.execute(
-        sa.text(
-            """
-            INSERT INTO slack_bot (name, enabled, bot_token, app_token)
-            SELECT 'Default Bot', true, '', ''
-            WHERE NOT EXISTS (SELECT 1 FROM slack_bot)
-            RETURNING id;
-            """
-        )
-    )

backend/alembic/versions/776b3bbe9092_remove_remaining_enums.py

@@ -9,8 +9,8 @@ from alembic import op
 import sqlalchemy as sa
 
 from danswer.db.models import IndexModelStatus
-from danswer.context.search.enums import RecencyBiasSetting
-from danswer.context.search.enums import SearchType
+from danswer.search.enums import RecencyBiasSetting
+from danswer.search.enums import SearchType
 
 # revision identifiers, used by Alembic.
 revision = "776b3bbe9092"

backend/alembic/versions/93560ba1b118_add_web_ui_option_to_slack_config.py

@@ -1,35 +0,0 @@
-"""add web ui option to slack config
-
-Revision ID: 93560ba1b118
-Revises: 6d562f86c78b
-Create Date: 2024-11-24 06:36:17.490612
-
-"""
-from alembic import op
-
-# revision identifiers, used by Alembic.
-revision = "93560ba1b118"
-down_revision = "6d562f86c78b"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Add show_continue_in_web_ui with default False to all existing channel_configs
-    op.execute(
-        """
-        UPDATE slack_channel_config
-        SET channel_config = channel_config || '{"show_continue_in_web_ui": false}'::jsonb
-        WHERE NOT channel_config ? 'show_continue_in_web_ui'
-        """
-    )
-
-
-def downgrade() -> None:
-    # Remove show_continue_in_web_ui from all channel_configs
-    op.execute(
-        """
-        UPDATE slack_channel_config
-        SET channel_config = channel_config - 'show_continue_in_web_ui'
-        """
-    )

backend/alembic/versions/949b4a92a401_remove_rt.py

@@ -7,7 +7,6 @@ Create Date: 2024-10-26 13:06:06.937969
 """
 from alembic import op
 from sqlalchemy.orm import Session
-from sqlalchemy import text
 
 # Import your models and constants
 from danswer.db.models import (
@@ -16,6 +15,7 @@ from danswer.db.models import (
     Credential,
     IndexAttempt,
 )
+from danswer.configs.constants import DocumentSource
 
 
 # revision identifiers, used by Alembic.
@@ -30,11 +30,13 @@ def upgrade() -> None:
     bind = op.get_bind()
     session = Session(bind=bind)
 
-    # Get connectors using raw SQL
-    result = bind.execute(
-        text("SELECT id FROM connector WHERE source = 'requesttracker'")
+    connectors_to_delete = (
+        session.query(Connector)
+        .filter(Connector.source == DocumentSource.REQUESTTRACKER)
+        .all()
     )
-    connector_ids = [row[0] for row in result]
+
+    connector_ids = [connector.id for connector in connectors_to_delete]
 
     if connector_ids:
         cc_pairs_to_delete = (

backend/alembic/versions/9cf5c00f72fe_add_creator_to_cc_pair.py

@@ -1,30 +0,0 @@
-"""add creator to cc pair
-
-Revision ID: 9cf5c00f72fe
-Revises: 26b931506ecb
-Create Date: 2024-11-12 15:16:42.682902
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "9cf5c00f72fe"
-down_revision = "26b931506ecb"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "connector_credential_pair",
-        sa.Column(
-            "creator_id",
-            sa.UUID(as_uuid=True),
-            nullable=True,
-        ),
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("connector_credential_pair", "creator_id")

backend/alembic/versions/9f696734098f_combine_search_and_chat.py

@@ -1,36 +0,0 @@
-"""Combine Search and Chat
-
-Revision ID: 9f696734098f
-Revises: a8c2065484e6
-Create Date: 2024-11-27 15:32:19.694972
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "9f696734098f"
-down_revision = "a8c2065484e6"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.alter_column("chat_session", "description", nullable=True)
-    op.drop_column("chat_session", "one_shot")
-    op.drop_column("slack_channel_config", "response_type")
-
-
-def downgrade() -> None:
-    op.execute("UPDATE chat_session SET description = '' WHERE description IS NULL")
-    op.alter_column("chat_session", "description", nullable=False)
-    op.add_column(
-        "chat_session",
-        sa.Column("one_shot", sa.Boolean(), nullable=False, server_default=sa.false()),
-    )
-    op.add_column(
-        "slack_channel_config",
-        sa.Column(
-            "response_type", sa.String(), nullable=False, server_default="citations"
-        ),
-    )

backend/alembic/versions/a8c2065484e6_add_auto_scroll_to_user_model.py

@@ -1,27 +0,0 @@
-"""add auto scroll to user model
-
-Revision ID: a8c2065484e6
-Revises: abe7378b8217
-Create Date: 2024-11-22 17:34:09.690295
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "a8c2065484e6"
-down_revision = "abe7378b8217"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "user",
-        sa.Column("auto_scroll", sa.Boolean(), nullable=True, server_default=None),
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("user", "auto_scroll")

backend/alembic/versions/abe7378b8217_add_indexing_trigger_to_cc_pair.py

@@ -1,30 +0,0 @@
-"""add indexing trigger to cc_pair
-
-Revision ID: abe7378b8217
-Revises: 6d562f86c78b
-Create Date: 2024-11-26 19:09:53.481171
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "abe7378b8217"
-down_revision = "93560ba1b118"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "connector_credential_pair",
-        sa.Column(
-            "indexing_trigger",
-            sa.Enum("UPDATE", "REINDEX", name="indexingmode", native_enum=False),
-            nullable=True,
-        ),
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("connector_credential_pair", "indexing_trigger")

backend/alembic/versions/b156fa702355_chat_reworked.py

@@ -288,15 +288,6 @@ 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/c99d76fcd298_add_nullable_to_persona_id_in_chat_.py

@@ -23,56 +23,6 @@ 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/alembic/versions/dfbe9e93d3c7_extended_role_for_non_web.py

@@ -1,42 +0,0 @@
-"""extended_role_for_non_web
-
-Revision ID: dfbe9e93d3c7
-Revises: 9cf5c00f72fe
-Create Date: 2024-11-16 07:54:18.727906
-
-"""
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "dfbe9e93d3c7"
-down_revision = "9cf5c00f72fe"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.execute(
-        """
-        UPDATE "user"
-        SET role = 'EXT_PERM_USER'
-        WHERE has_web_login = false
-    """
-    )
-    op.drop_column("user", "has_web_login")
-
-
-def downgrade() -> None:
-    op.add_column(
-        "user",
-        sa.Column("has_web_login", sa.Boolean(), nullable=False, server_default="true"),
-    )
-
-    op.execute(
-        """
-        UPDATE "user"
-        SET has_web_login = false,
-            role = 'BASIC'
-        WHERE role IN ('SLACK_USER', 'EXT_PERM_USER')
-    """
-    )

backend/alembic_tenants/env.py

@@ -1,6 +1,5 @@
 import asyncio
 from logging.config import fileConfig
-from typing import Literal
 
 from sqlalchemy import pool
 from sqlalchemy.engine import Connection
@@ -38,15 +37,8 @@ EXCLUDE_TABLES = {"kombu_queue", "kombu_message"}
 
 def include_object(
     object: SchemaItem,
-    name: str | None,
-    type_: Literal[
-        "schema",
-        "table",
-        "column",
-        "index",
-        "unique_constraint",
-        "foreign_key_constraint",
-    ],
+    name: str,
+    type_: str,
     reflected: bool,
     compare_to: SchemaItem | None,
 ) -> bool:

backend/danswer/access/models.py

@@ -16,46 +16,6 @@ class ExternalAccess:
     is_public: bool
 
 
-@dataclass(frozen=True)
-class DocExternalAccess:
-    """
-    This is just a class to wrap the external access and the document ID
-    together. It's used for syncing document permissions to Redis.
-    """
-
-    external_access: ExternalAccess
-    # The document ID
-    doc_id: str
-
-    def to_dict(self) -> dict:
-        return {
-            "external_access": {
-                "external_user_emails": list(self.external_access.external_user_emails),
-                "external_user_group_ids": list(
-                    self.external_access.external_user_group_ids
-                ),
-                "is_public": self.external_access.is_public,
-            },
-            "doc_id": self.doc_id,
-        }
-
-    @classmethod
-    def from_dict(cls, data: dict) -> "DocExternalAccess":
-        external_access = ExternalAccess(
-            external_user_emails=set(
-                data["external_access"].get("external_user_emails", [])
-            ),
-            external_user_group_ids=set(
-                data["external_access"].get("external_user_group_ids", [])
-            ),
-            is_public=data["external_access"]["is_public"],
-        )
-        return cls(
-            external_access=external_access,
-            doc_id=data["doc_id"],
-        )
-
-
 @dataclass(frozen=True)
 class DocumentAccess(ExternalAccess):
     # User emails for Danswer users, None indicates admin

backend/danswer/agent_search/core_qa_graph/edges.py

@@ -1,42 +0,0 @@
-from collections.abc import Hashable
-from typing import Union
-
-from langgraph.types import Send
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.primary_graph.states import RetrieverState
-from danswer.agent_search.primary_graph.states import VerifierState
-
-
-def sub_continue_to_verifier(state: BaseQAState) -> Union[Hashable, list[Hashable]]:
-    # Routes each de-douped retrieved doc to the verifier step - in parallel
-    # Notice the 'Send()' API that takes care of the parallelization
-
-    return [
-        Send(
-            "sub_verifier",
-            VerifierState(
-                document=doc,
-                #question=state["original_question"],
-                question=state["sub_question_str"],
-                graph_start_time=state["graph_start_time"],
-            ),
-        )
-        for doc in state["sub_question_deduped_retrieval_docs"]
-    ]
-
-
-def sub_continue_to_retrieval(state: BaseQAState) -> Union[Hashable, list[Hashable]]:
-    # Routes re-written queries to the (parallel) retrieval steps
-    # Notice the 'Send()' API that takes care of the parallelization
-    rewritten_queries = state["sub_question_search_queries"].rewritten_queries + [state["sub_question_str"]]
-    return [
-        Send(
-            "sub_custom_retrieve",
-            RetrieverState(
-                rewritten_query=query,
-                graph_start_time=state["graph_start_time"],
-            ),
-        )
-        for query in rewritten_queries
-    ]

backend/danswer/agent_search/core_qa_graph/graph_builder.py

@@ -1,132 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from danswer.agent_search.core_qa_graph.edges import sub_continue_to_retrieval
-from danswer.agent_search.core_qa_graph.edges import sub_continue_to_verifier
-from danswer.agent_search.core_qa_graph.nodes.combine_retrieved_docs import (
-    sub_combine_retrieved_docs,
-)
-from danswer.agent_search.core_qa_graph.nodes.custom_retrieve import (
-    sub_custom_retrieve,
-)
-from danswer.agent_search.core_qa_graph.nodes.dummy import sub_dummy
-from danswer.agent_search.core_qa_graph.nodes.final_format import (
-    sub_final_format,
-)
-from danswer.agent_search.core_qa_graph.nodes.generate import sub_generate
-from danswer.agent_search.core_qa_graph.nodes.qa_check import sub_qa_check
-from danswer.agent_search.core_qa_graph.nodes.rewrite import sub_rewrite
-from danswer.agent_search.core_qa_graph.nodes.verifier import sub_verifier
-from danswer.agent_search.core_qa_graph.states import BaseQAOutputState
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.core_qa_graph.states import CoreQAInputState
-
-
-def build_core_qa_graph() -> StateGraph:
-    sub_answers_initial = StateGraph(
-        state_schema=BaseQAState,
-        output=BaseQAOutputState,
-    )
-
-    ### Add nodes ###
-    sub_answers_initial.add_node(node="sub_dummy", action=sub_dummy)
-    sub_answers_initial.add_node(node="sub_rewrite", action=sub_rewrite)
-    sub_answers_initial.add_node(
-        node="sub_custom_retrieve",
-        action=sub_custom_retrieve,
-    )
-    sub_answers_initial.add_node(
-        node="sub_combine_retrieved_docs",
-        action=sub_combine_retrieved_docs,
-    )
-    sub_answers_initial.add_node(
-        node="sub_verifier",
-        action=sub_verifier,
-    )
-    sub_answers_initial.add_node(
-        node="sub_generate",
-        action=sub_generate,
-    )
-    sub_answers_initial.add_node(
-        node="sub_qa_check",
-        action=sub_qa_check,
-    )
-    sub_answers_initial.add_node(
-        node="sub_final_format",
-        action=sub_final_format,
-    )
-
-    ### Add edges ###
-    sub_answers_initial.add_edge(START, "sub_dummy")
-    sub_answers_initial.add_edge("sub_dummy", "sub_rewrite")
-
-    sub_answers_initial.add_conditional_edges(
-        source="sub_rewrite",
-        path=sub_continue_to_retrieval,
-    )
-
-    sub_answers_initial.add_edge(
-        start_key="sub_custom_retrieve",
-        end_key="sub_combine_retrieved_docs",
-    )
-
-    sub_answers_initial.add_conditional_edges(
-        source="sub_combine_retrieved_docs",
-        path=sub_continue_to_verifier,
-        path_map=["sub_verifier"],
-    )
-
-    sub_answers_initial.add_edge(
-        start_key="sub_verifier",
-        end_key="sub_generate",
-    )
-
-    sub_answers_initial.add_edge(
-        start_key="sub_generate",
-        end_key="sub_qa_check",
-    )
-
-    sub_answers_initial.add_edge(
-        start_key="sub_qa_check",
-        end_key="sub_final_format",
-    )
-
-    sub_answers_initial.add_edge(
-        start_key="sub_final_format",
-        end_key=END,
-    )
-    # sub_answers_graph = sub_answers_initial.compile()
-    return sub_answers_initial
-
-
-if __name__ == "__main__":
-    # q = "Whose music is kind of hard to easily enjoy?"
-    # q = "What is voice leading?"
-    # q = "What are the types of motions in music?"
-    # q = "What are key elements of music theory?"
-    # q = "How can I best understand music theory using voice leading?"
-    q = "What makes good music?"
-    # q = "types of motions in music"
-    # q = "What is the relationship between music and physics?"
-    # q = "Can you compare various grunge styles?"
-    # q = "Why is quantum gravity so hard?"
-
-    inputs = CoreQAInputState(
-        original_question=q,
-        sub_question_str=q,
-    )
-    sub_answers_graph = build_core_qa_graph()
-    compiled_sub_answers = sub_answers_graph.compile()
-    output = compiled_sub_answers.invoke(inputs)
-    print("\nOUTPUT:")
-    print(output.keys())
-    for key, value in output.items():
-        if key in [
-            "sub_question_answer",
-            "sub_question_str",
-            "sub_qas",
-            "initial_sub_qas",
-            "sub_question_answer",
-        ]:
-            print(f"{key}: {value}")

backend/danswer/agent_search/core_qa_graph/nodes/combine_retrieved_docs.py

@@ -1,36 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.context.search.models import InferenceSection
-
-
-def sub_combine_retrieved_docs(state: BaseQAState) -> dict[str, Any]:
-    """
-    Dedupe the retrieved docs.
-    """
-    node_start_time = datetime.now()
-
-    sub_question_base_retrieval_docs = state["sub_question_base_retrieval_docs"]
-
-    print(f"Number of docs from steps: {len(sub_question_base_retrieval_docs)}")
-    dedupe_docs: list[InferenceSection] = []
-    for base_retrieval_doc in sub_question_base_retrieval_docs:
-        if not any(
-            base_retrieval_doc.center_chunk.chunk_id == doc.center_chunk.chunk_id
-            for doc in dedupe_docs
-        ):
-            dedupe_docs.append(base_retrieval_doc)
-
-    print(f"Number of deduped docs: {len(dedupe_docs)}")
-    
-
-    return {
-        "sub_question_deduped_retrieval_docs": dedupe_docs,
-        "log_messages": generate_log_message(
-            message="sub - combine_retrieved_docs (dedupe)",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/nodes/custom_retrieve.py

@@ -1,66 +0,0 @@
-import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import RetrieverState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.context.search.models import InferenceSection
-from danswer.context.search.models import SearchRequest
-from danswer.context.search.pipeline import SearchPipeline
-from danswer.db.engine import get_session_context_manager
-from danswer.llm.factory import get_default_llms
-
-
-def sub_custom_retrieve(state: RetrieverState) -> dict[str, Any]:
-    """
-    Retrieve documents
-
-    Args:
-        state (dict): The current graph state
-
-    Returns:
-        state (dict): New key added to state, documents, that contains retrieved documents
-    """
-    print("---RETRIEVE SUB---")
-
-    node_start_time = datetime.datetime.now()
-
-    rewritten_query = state["rewritten_query"]
-
-    # Retrieval
-    # TODO: add the actual retrieval, probably from search_tool.run()
-    documents: list[InferenceSection] = []
-    llm, fast_llm = get_default_llms()
-    with get_session_context_manager() as db_session:
-        documents = SearchPipeline(
-            search_request=SearchRequest(
-                query=rewritten_query,
-            ),
-            user=None,
-            llm=llm,
-            fast_llm=fast_llm,
-            db_session=db_session,
-        )
-        
-        reranked_docs = documents.reranked_sections
-
-        # initial metric to measure fit TODO: implement metric properly
-
-        top_1_score = reranked_docs[0].center_chunk.score
-        top_5_score = sum([doc.center_chunk.score for doc in reranked_docs[:5]]) / 5
-        top_10_score = sum([doc.center_chunk.score for doc in reranked_docs[:10]]) / 10
-
-        fit_score = 1/3 * (top_1_score + top_5_score + top_10_score)
-
-        chunk_ids = {'query': rewritten_query, 
-                     'chunk_ids': [doc.center_chunk.chunk_id for doc in reranked_docs]}
-
-
-    return {
-        "sub_question_base_retrieval_docs": reranked_docs,
-        "sub_chunk_ids": [chunk_ids],
-        "log_messages": generate_log_message(
-            message=f"sub - custom_retrieve, fit_score: {fit_score}",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/nodes/dummy.py

@@ -1,24 +0,0 @@
-import datetime
-from typing import Any
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_dummy(state: BaseQAState) -> dict[str, Any]:
-    """
-    Dummy step
-    """
-
-    print("---Sub Dummy---")
-
-    node_start_time = datetime.datetime.now()
-
-    return {
-        "graph_start_time": node_start_time,
-        "log_messages": generate_log_message(
-            message="sub - dummy",
-            node_start_time=node_start_time,
-            graph_start_time=node_start_time,
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/nodes/final_format.py

@@ -1,22 +0,0 @@
-from typing import Any
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-
-
-def sub_final_format(state: BaseQAState) -> dict[str, Any]:
-    """
-    Create the final output for the QA subgraph
-    """
-
-    print("---BASE FINAL FORMAT---")
-
-    return {
-        "sub_qas": [
-            {
-                "sub_question": state["sub_question_str"],
-                "sub_answer": state["sub_question_answer"],
-                "sub_answer_check": state["sub_question_answer_check"],
-            }
-        ],
-        "log_messages": state["log_messages"],
-    }

backend/danswer/agent_search/core_qa_graph/nodes/generate.py

@@ -1,91 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.shared_graph_utils.prompts import BASE_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.llm.factory import get_default_llms
-
-
-def sub_generate(state: BaseQAState) -> dict[str, Any]:
-    """
-    Generate answer
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-         dict: The updated state with re-phrased question
-    """
-    print("---GENERATE---")
-    
-    # Create sub-query results
-
-    verified_chunks = [chunk.center_chunk.chunk_id for chunk in state["sub_question_verified_retrieval_docs"]]
-    result_dict = {}
-
-    chunk_id_dicts = state["sub_chunk_ids"]
-    expanded_chunks = []
-    original_chunks = []
-    
-    for chunk_id_dict in chunk_id_dicts:
-        sub_question = chunk_id_dict['query']
-        verified_sq_chunks = [chunk_id for chunk_id in chunk_id_dict['chunk_ids'] if chunk_id in verified_chunks]
-
-        if sub_question != state["original_question"]:
-            expanded_chunks += verified_sq_chunks
-        else:
-            result_dict['ORIGINAL'] = len(verified_sq_chunks)
-            original_chunks += verified_sq_chunks
-        result_dict[sub_question[:30]] = len(verified_sq_chunks)
-    
-    expansion_chunks = set(expanded_chunks)
-    num_expansion_chunks = sum([1 for chunk_id in expansion_chunks if chunk_id in verified_chunks])
-    num_original_relevant_chunks = len(original_chunks)
-    num_missed_relevant_chunks = sum([1 for chunk_id in original_chunks if chunk_id not in expansion_chunks])
-    num_gained_relevant_chunks = sum([1 for chunk_id in expansion_chunks if chunk_id not in original_chunks])
-    result_dict['expansion_chunks'] = num_expansion_chunks
-
-
-        
-    print(result_dict)
-
-    node_start_time = datetime.now()
-
-    question = state["sub_question_str"]
-    docs = state["sub_question_verified_retrieval_docs"]
-
-    print(f"Number of verified retrieval docs: {len(docs)}")
-    
-    # Only take the top 10 docs. 
-    # TODO: Make this dynamic or use config param?
-    top_10_docs = docs[-10:]
-
-    msg = [
-        HumanMessage(
-            content=BASE_RAG_PROMPT.format(question=question, context=format_docs(top_10_docs))
-        )
-    ]
-
-    # Grader
-    _, fast_llm = get_default_llms()
-    response = list(
-        fast_llm.stream(
-            prompt=msg,
-            # structured_response_format=None,
-        )
-    )
-
-    answer_str = merge_message_runs(response, chunk_separator="")[0].content
-    return {
-        "sub_question_answer": answer_str,
-        "log_messages": generate_log_message(
-            message="base - generate",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/nodes/qa_check.py

@@ -1,51 +0,0 @@
-import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.shared_graph_utils.prompts import BASE_CHECK_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.llm.factory import get_default_llms
-
-
-def sub_qa_check(state: BaseQAState) -> dict[str, Any]:
-    """
-    Check if the sub-question answer is satisfactory.
-
-    Args:
-        state: The current SubQAState containing the sub-question and its answer
-
-    Returns:
-        dict containing the check result and log message
-    """
-    node_start_time = datetime.datetime.now()
-
-    msg = [
-        HumanMessage(
-            content=BASE_CHECK_PROMPT.format(
-                question=state["sub_question_str"],
-                base_answer=state["sub_question_answer"],
-            )
-        )
-    ]
-
-    _, fast_llm = get_default_llms()
-    response = list(
-        fast_llm.stream(
-            prompt=msg,
-            # structured_response_format=None,
-        )
-    )
-
-    response_str = merge_message_runs(response, chunk_separator="")[0].content
-
-    return {
-        "sub_question_answer_check": response_str,
-        "base_answer_messages": generate_log_message(
-            message="sub - qa_check",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/nodes/rewrite.py

@@ -1,74 +0,0 @@
-import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.shared_graph_utils.models import RewrittenQueries
-from danswer.agent_search.shared_graph_utils.prompts import (
-    REWRITE_PROMPT_MULTI_ORIGINAL,
-)
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.llm.factory import get_default_llms
-
-
-def sub_rewrite(state: BaseQAState) -> dict[str, Any]:
-    """
-    Transform the initial question into more suitable search queries.
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: The updated state with re-phrased question
-    """
-
-    print("---SUB TRANSFORM QUERY---")
-
-    node_start_time = datetime.datetime.now()
-
-    # messages = state["base_answer_messages"]
-    question = state["sub_question_str"]
-
-    msg = [
-        HumanMessage(
-            content=REWRITE_PROMPT_MULTI_ORIGINAL.format(question=question),
-        )
-    ]
-
-    """
-    msg = [
-        HumanMessage(
-            content=REWRITE_PROMPT_MULTI.format(question=question),
-        )
-    ]
-    """
-
-    _, fast_llm = get_default_llms()
-    llm_response_list = list(
-        fast_llm.stream(
-            prompt=msg,
-            # structured_response_format={"type": "json_object", "schema": RewrittenQueries.model_json_schema()},
-            # structured_response_format=RewrittenQueries.model_json_schema(),
-        )
-    )
-    llm_response = merge_message_runs(llm_response_list, chunk_separator="")[0].content
-
-    print(f"llm_response: {llm_response}")
-
-    rewritten_queries = llm_response.split("--")
-    # rewritten_queries = [llm_response.split("\n")[0]]
-
-    print(f"rewritten_queries: {rewritten_queries}")
-
-    rewritten_queries = RewrittenQueries(rewritten_queries=rewritten_queries)
-
-    return {
-        "sub_question_search_queries": rewritten_queries,
-        "log_messages": generate_log_message(
-            message="sub - rewrite",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/nodes/verifier.py

@@ -1,64 +0,0 @@
-import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.primary_graph.states import VerifierState
-from danswer.agent_search.shared_graph_utils.models import BinaryDecision
-from danswer.agent_search.shared_graph_utils.prompts import VERIFIER_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.llm.factory import get_default_llms
-
-
-def sub_verifier(state: VerifierState) -> dict[str, Any]:
-    """
-    Check whether the document is relevant for the original user question
-
-    Args:
-        state (VerifierState): The current state
-
-    Returns:
-        dict: ict: The updated state with the final decision
-    """
-
-    # print("---VERIFY QUTPUT---")
-    node_start_time = datetime.datetime.now()
-
-    question = state["question"]
-    document_content = state["document"].combined_content
-
-    msg = [
-        HumanMessage(
-            content=VERIFIER_PROMPT.format(
-                question=question, document_content=document_content
-            )
-        )
-    ]
-
-    # Grader
-    llm, fast_llm = get_default_llms()
-    response = list(
-        llm.stream(
-            prompt=msg,
-            # structured_response_format=BinaryDecision.model_json_schema(),
-        )
-    )
-
-    response_string = merge_message_runs(response, chunk_separator="")[0].content
-    # Convert string response to proper dictionary format
-    decision_dict = {"decision": response_string.lower()}
-    formatted_response = BinaryDecision.model_validate(decision_dict)
-
-    print(f"Verification end time: {datetime.datetime.now()}")
-
-    return {
-        "sub_question_verified_retrieval_docs": [state["document"]]
-        if formatted_response.decision == "yes"
-        else [],
-        "log_messages": generate_log_message(
-            message=f"sub - verifier: {formatted_response.decision}",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/core_qa_graph/states.py

@@ -1,90 +0,0 @@
-import operator
-from collections.abc import Sequence
-from datetime import datetime
-from typing import Annotated
-from typing import TypedDict
-
-from langchain_core.messages import BaseMessage
-from langgraph.graph.message import add_messages
-
-from danswer.agent_search.shared_graph_utils.models import RewrittenQueries
-from danswer.context.search.models import InferenceSection
-from danswer.llm.interfaces import LLM
-
-
-class SubQuestionRetrieverState(TypedDict):
-    # The state for the parallel Retrievers. They each need to see only one query
-    sub_question_rewritten_query: str
-
-
-class SubQuestionVerifierState(TypedDict):
-    # The state for the parallel verification step.  Each node execution need to see only one question/doc pair
-    sub_question_document: InferenceSection
-    sub_question: str
-
-
-class CoreQAInputState(TypedDict):
-    sub_question_str: str
-    original_question: str
-
-
-class BaseQAState(TypedDict):
-    # The 'core SubQuestion'  state.
-    original_question: str
-    graph_start_time: datetime
-    # start time for parallel initial sub-questionn thread
-    sub_query_start_time: datetime
-    sub_question_rewritten_queries: list[str]
-    sub_question_str: str
-    sub_question_search_queries: RewrittenQueries
-    sub_question_nr: int
-    sub_chunk_ids: Annotated[Sequence[dict], operator.add]
-    sub_question_base_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_deduped_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_verified_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_reranked_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_top_chunks: Annotated[Sequence[dict], operator.add]
-    sub_question_answer: str
-    sub_question_answer_check: str
-    log_messages: Annotated[Sequence[BaseMessage], add_messages]
-    sub_qas: Annotated[Sequence[dict], operator.add]
-    # Answers sent back to core
-    initial_sub_qas: Annotated[Sequence[dict], operator.add]
-    primary_llm: LLM
-    fast_llm: LLM
-
-
-class BaseQAOutputState(TypedDict):
-    # The 'SubQuestion'  output state. Removes all the intermediate states
-    sub_question_rewritten_queries: list[str]
-    sub_question_str: str
-    sub_question_search_queries: list[str]
-    sub_question_nr: int
-    # Answers sent back to core
-    sub_qas: Annotated[Sequence[dict], operator.add]
-    # Answers sent back to core
-    initial_sub_qas: Annotated[Sequence[dict], operator.add]
-    sub_question_base_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_deduped_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_verified_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_reranked_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_top_chunks: Annotated[Sequence[dict], operator.add]
-    sub_question_answer: str
-    sub_question_answer_check: str
-    log_messages: Annotated[Sequence[BaseMessage], add_messages]

backend/danswer/agent_search/deep_qa_graph/edges.py

@@ -1,46 +0,0 @@
-from collections.abc import Hashable
-from typing import Union
-
-from langgraph.types import Send
-
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.primary_graph.states import RetrieverState
-from danswer.agent_search.primary_graph.states import VerifierState
-
-
-def sub_continue_to_verifier(state: ResearchQAState) -> Union[Hashable, list[Hashable]]:
-    # Routes each de-douped retrieved doc to the verifier step - in parallel
-    # Notice the 'Send()' API that takes care of the parallelization
-
-    return [
-        Send(
-            "sub_verifier",
-            VerifierState(
-                document=doc,
-                question=state["sub_question"],
-                primary_llm=state["primary_llm"],
-                fast_llm=state["fast_llm"],
-                graph_start_time=state["graph_start_time"],
-            ),
-        )
-        for doc in state["sub_question_base_retrieval_docs"]
-    ]
-
-
-def sub_continue_to_retrieval(
-    state: ResearchQAState,
-) -> Union[Hashable, list[Hashable]]:
-    # Routes re-written queries to the (parallel) retrieval steps
-    # Notice the 'Send()' API that takes care of the parallelization
-    return [
-        Send(
-            "sub_custom_retrieve",
-            RetrieverState(
-                rewritten_query=query,
-                primary_llm=state["primary_llm"],
-                fast_llm=state["fast_llm"],
-                graph_start_time=state["graph_start_time"],
-            ),
-        )
-        for query in state["sub_question_rewritten_queries"]
-    ]

backend/danswer/agent_search/deep_qa_graph/graph_builder.py

@@ -1,93 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from danswer.agent_search.deep_qa_graph.edges import sub_continue_to_retrieval
-from danswer.agent_search.deep_qa_graph.edges import sub_continue_to_verifier
-from danswer.agent_search.deep_qa_graph.nodes.combine_retrieved_docs import (
-    sub_combine_retrieved_docs,
-)
-from danswer.agent_search.deep_qa_graph.nodes.custom_retrieve import sub_custom_retrieve
-from danswer.agent_search.deep_qa_graph.nodes.dummy import sub_dummy
-from danswer.agent_search.deep_qa_graph.nodes.final_format import sub_final_format
-from danswer.agent_search.deep_qa_graph.nodes.generate import sub_generate
-from danswer.agent_search.deep_qa_graph.nodes.qa_check import sub_qa_check
-from danswer.agent_search.deep_qa_graph.nodes.verifier import sub_verifier
-from danswer.agent_search.deep_qa_graph.states import ResearchQAOutputState
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-
-
-def build_deep_qa_graph() -> StateGraph:
-    # Define the nodes we will cycle between
-    sub_answers = StateGraph(state_schema=ResearchQAState, output=ResearchQAOutputState)
-
-    ### Add Nodes ###
-
-    # Dummy node for initial processing
-    sub_answers.add_node(node="sub_dummy", action=sub_dummy)
-
-    # The retrieval step
-    sub_answers.add_node(node="sub_custom_retrieve", action=sub_custom_retrieve)
-
-    # The dedupe step
-    sub_answers.add_node(
-        node="sub_combine_retrieved_docs", action=sub_combine_retrieved_docs
-    )
-
-    # Verifying retrieved information
-    sub_answers.add_node(node="sub_verifier", action=sub_verifier)
-
-    # Generating the response
-    sub_answers.add_node(node="sub_generate", action=sub_generate)
-
-    # Checking the quality of the answer
-    sub_answers.add_node(node="sub_qa_check", action=sub_qa_check)
-
-    # Final formatting of the response
-    sub_answers.add_node(node="sub_final_format", action=sub_final_format)
-
-    ### Add Edges ###
-
-    # Generate multiple sub-questions
-    sub_answers.add_edge(start_key=START, end_key="sub_rewrite")
-
-    # For each sub-question, perform a retrieval in parallel
-    sub_answers.add_conditional_edges(
-        source="sub_rewrite",
-        path=sub_continue_to_retrieval,
-        path_map=["sub_custom_retrieve"],
-    )
-
-    # Combine the retrieved docs for each sub-question from the parallel retrievals
-    sub_answers.add_edge(
-        start_key="sub_custom_retrieve", end_key="sub_combine_retrieved_docs"
-    )
-
-    # Go over all of the combined retrieved docs and verify them against the original question
-    sub_answers.add_conditional_edges(
-        source="sub_combine_retrieved_docs",
-        path=sub_continue_to_verifier,
-        path_map=["sub_verifier"],
-    )
-
-    # Generate an answer for each verified retrieved doc
-    sub_answers.add_edge(start_key="sub_verifier", end_key="sub_generate")
-
-    # Check the quality of the answer
-    sub_answers.add_edge(start_key="sub_generate", end_key="sub_qa_check")
-
-    sub_answers.add_edge(start_key="sub_qa_check", end_key="sub_final_format")
-
-    sub_answers.add_edge(start_key="sub_final_format", end_key=END)
-
-    return sub_answers
-
-
-if __name__ == "__main__":
-    # TODO: add the actual question
-    inputs = {"sub_question": "Whose music is kind of hard to easily enjoy?"}
-    sub_answers_graph = build_deep_qa_graph()
-    compiled_sub_answers = sub_answers_graph.compile()
-    output = compiled_sub_answers.invoke(inputs)
-    print("\nOUTPUT:")
-    print(output)

backend/danswer/agent_search/deep_qa_graph/nodes/combine_retrieved_docs.py

@@ -1,31 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_combine_retrieved_docs(state: ResearchQAState) -> dict[str, Any]:
-    """
-    Dedupe the retrieved docs.
-    """
-    node_start_time = datetime.now()
-
-    sub_question_base_retrieval_docs = state["sub_question_base_retrieval_docs"]
-
-    print(f"Number of docs from steps: {len(sub_question_base_retrieval_docs)}")
-    dedupe_docs = []
-    for base_retrieval_doc in sub_question_base_retrieval_docs:
-        if base_retrieval_doc not in dedupe_docs:
-            dedupe_docs.append(base_retrieval_doc)
-
-    print(f"Number of deduped docs: {len(dedupe_docs)}")
-
-    return {
-        "sub_question_deduped_retrieval_docs": dedupe_docs,
-        "log_messages": generate_log_message(
-            message="sub - combine_retrieved_docs (dedupe)",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/custom_retrieve.py

@@ -1,33 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import RetrieverState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.context.search.models import InferenceSection
-
-
-def sub_custom_retrieve(state: RetrieverState) -> dict[str, Any]:
-    """
-    Retrieve documents
-
-    Args:
-        state (dict): The current graph state
-
-    Returns:
-        state (dict): New key added to state, documents, that contains retrieved documents
-    """
-    print("---RETRIEVE SUB---")
-    node_start_time = datetime.now()
-
-    # Retrieval
-    # TODO: add the actual retrieval, probably from search_tool.run()
-    documents: list[InferenceSection] = []
-
-    return {
-        "sub_question_base_retrieval_docs": documents,
-        "log_messages": generate_log_message(
-            message="sub - custom_retrieve",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/dummy.py

@@ -1,21 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_dummy(state: BaseQAState) -> dict[str, Any]:
-    """
-    Dummy step
-    """
-
-    print("---Sub Dummy---")
-
-    return {
-        "log_messages": generate_log_message(
-            message="sub - dummy",
-            node_start_time=datetime.now(),
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/final_format.py

@@ -1,31 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_final_format(state: ResearchQAState) -> dict[str, Any]:
-    """
-    Create the final output for the QA subgraph
-    """
-
-    print("---SUB  FINAL FORMAT---")
-    node_start_time = datetime.now()
-
-    return {
-        # TODO: Type this
-        "sub_qas": [
-            {
-                "sub_question": state["sub_question"],
-                "sub_answer": state["sub_question_answer"],
-                "sub_question_nr": state["sub_question_nr"],
-                "sub_answer_check": state["sub_question_answer_check"],
-            }
-        ],
-        "log_messages": generate_log_message(
-            message="sub - final format",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/generate.py

@@ -1,56 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.shared_graph_utils.prompts import BASE_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_generate(state: ResearchQAState) -> dict[str, Any]:
-    """
-    Generate answer
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-         dict: The updated state with re-phrased question
-    """
-    print("---SUB GENERATE---")
-    node_start_time = datetime.now()
-
-    question = state["sub_question"]
-    docs = state["sub_question_verified_retrieval_docs"]
-
-    print(f"Number of verified retrieval docs for sub-question: {len(docs)}")
-
-    msg = [
-        HumanMessage(
-            content=BASE_RAG_PROMPT.format(question=question, context=format_docs(docs))
-        )
-    ]
-
-    # Grader
-    if len(docs) > 0:
-        model = state["fast_llm"]
-        response = list(
-            model.stream(
-                prompt=msg,
-            )
-        )
-        response_str = merge_message_runs(response, chunk_separator="")[0].content
-    else:
-        response_str = ""
-
-    return {
-        "sub_question_answer": response_str,
-        "log_messages": generate_log_message(
-            message="sub - generate",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/qa_check.py

@@ -1,57 +0,0 @@
-import json
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.deep_qa_graph.prompts import SUB_CHECK_PROMPT
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.shared_graph_utils.models import BinaryDecision
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_qa_check(state: ResearchQAState) -> dict[str, Any]:
-    """
-    Check whether the final output satisfies the original user question
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: The updated state with the final decision
-    """
-
-    print("---CHECK SUB QUTPUT---")
-    node_start_time = datetime.now()
-
-    sub_answer = state["sub_question_answer"]
-    sub_question = state["sub_question"]
-
-    msg = [
-        HumanMessage(
-            content=SUB_CHECK_PROMPT.format(
-                sub_question=sub_question, sub_answer=sub_answer
-            )
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = list(
-        model.stream(
-            prompt=msg,
-            structured_response_format=BinaryDecision.model_json_schema(),
-        )
-    )
-
-    raw_response = json.loads(response[0].pretty_repr())
-    formatted_response = BinaryDecision.model_validate(raw_response)
-
-    return {
-        "sub_question_answer_check": formatted_response.decision,
-        "log_messages": generate_log_message(
-            message=f"sub - qa check: {formatted_response.decision}",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/rewrite.py

@@ -1,64 +0,0 @@
-import json
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.shared_graph_utils.models import RewrittenQueries
-from danswer.agent_search.shared_graph_utils.prompts import REWRITE_PROMPT_MULTI
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.llm.interfaces import LLM
-
-
-def sub_rewrite(state: ResearchQAState) -> dict[str, Any]:
-    """
-    Transform the initial question into more suitable search queries.
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: The updated state with re-phrased question
-    """
-
-    print("---SUB TRANSFORM QUERY---")
-    node_start_time = datetime.now()
-
-    question = state["sub_question"]
-
-    msg = [
-        HumanMessage(
-            content=REWRITE_PROMPT_MULTI.format(question=question),
-        )
-    ]
-    fast_llm: LLM = state["fast_llm"]
-    llm_response = list(
-        fast_llm.stream(
-            prompt=msg,
-            structured_response_format=RewrittenQueries.model_json_schema(),
-        )
-    )
-
-    # Get the rewritten queries in a defined format
-    rewritten_queries: RewrittenQueries = json.loads(llm_response[0].pretty_repr())
-
-    print(f"rewritten_queries: {rewritten_queries}")
-
-    rewritten_queries = RewrittenQueries(
-        rewritten_queries=[
-            "music hard to listen to",
-            "Music that is not fun or pleasant",
-        ]
-    )
-
-    print(f"hardcoded rewritten_queries: {rewritten_queries}")
-
-    return {
-        "sub_question_rewritten_queries": rewritten_queries,
-        "log_messages": generate_log_message(
-            message="sub - rewrite",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/nodes/verifier.py

@@ -1,59 +0,0 @@
-import json
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.states import VerifierState
-from danswer.agent_search.shared_graph_utils.models import BinaryDecision
-from danswer.agent_search.shared_graph_utils.prompts import VERIFIER_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_verifier(state: VerifierState) -> dict[str, Any]:
-    """
-    Check whether the document is relevant for the original user question
-
-    Args:
-        state (VerifierState): The current state
-
-    Returns:
-        dict: ict: The updated state with the final decision
-    """
-
-    print("---SUB VERIFY QUTPUT---")
-    node_start_time = datetime.now()
-
-    question = state["question"]
-    document_content = state["document"].combined_content
-
-    msg = [
-        HumanMessage(
-            content=VERIFIER_PROMPT.format(
-                question=question, document_content=document_content
-            )
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = list(
-        model.stream(
-            prompt=msg,
-            structured_response_format=BinaryDecision.model_json_schema(),
-        )
-    )
-
-    raw_response = json.loads(response[0].pretty_repr())
-    formatted_response = BinaryDecision.model_validate(raw_response)
-
-    return {
-        "deduped_retrieval_docs": [state["document"]]
-        if formatted_response.decision == "yes"
-        else [],
-        "log_messages": generate_log_message(
-            message=f"core - verifier: {formatted_response.decision}",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/deep_qa_graph/prompts.py

@@ -1,13 +0,0 @@
-SUB_CHECK_PROMPT = """ \n
-    Please check whether the suggested answer seems to address the original question.
-
-    Please only answer with 'yes' or 'no' \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Here is the proposed answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-    Please answer with yes or no:"""

backend/danswer/agent_search/deep_qa_graph/states.py

@@ -1,64 +0,0 @@
-import operator
-from collections.abc import Sequence
-from datetime import datetime
-from typing import Annotated
-from typing import TypedDict
-
-from langchain_core.messages import BaseMessage
-from langgraph.graph.message import add_messages
-
-from danswer.context.search.models import InferenceSection
-from danswer.llm.interfaces import LLM
-
-
-class ResearchQAState(TypedDict):
-    # The 'core SubQuestion'  state.
-    original_question: str
-    graph_start_time: datetime
-    sub_question_rewritten_queries: list[str]
-    sub_question: str
-    sub_question_nr: int
-    sub_question_base_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_deduped_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_verified_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_reranked_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_top_chunks: Annotated[Sequence[dict], operator.add]
-    sub_question_answer: str
-    sub_question_answer_check: str
-    log_messages: Annotated[Sequence[BaseMessage], add_messages]
-    sub_qas: Annotated[Sequence[dict], operator.add]
-    primary_llm: LLM
-    fast_llm: LLM
-
-
-class ResearchQAOutputState(TypedDict):
-    # The 'SubQuestion'  output state. Removes all the intermediate states
-    sub_question_rewritten_queries: list[str]
-    sub_question: str
-    sub_question_nr: int
-    # Answers sent back to core
-    sub_qas: Annotated[Sequence[dict], operator.add]
-    sub_question_base_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_deduped_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_verified_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_reranked_retrieval_docs: Annotated[
-        Sequence[InferenceSection], operator.add
-    ]
-    sub_question_top_chunks: Annotated[Sequence[dict], operator.add]
-    sub_question_answer: str
-    sub_question_answer_check: str
-    log_messages: Annotated[Sequence[BaseMessage], add_messages]

backend/danswer/agent_search/primary_graph/edges.py

@@ -1,75 +0,0 @@
-from collections.abc import Hashable
-from typing import Union
-
-from langchain_core.messages import HumanMessage
-from langgraph.types import Send
-
-from danswer.agent_search.core_qa_graph.states import BaseQAState
-from danswer.agent_search.deep_qa_graph.states import ResearchQAState
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.prompts import BASE_CHECK_PROMPT
-
-
-def continue_to_initial_sub_questions(
-    state: QAState,
-) -> Union[Hashable, list[Hashable]]:
-    # Routes re-written queries to the (parallel) retrieval steps
-    # Notice the 'Send()' API that takes care of the parallelization
-    return [
-        Send(
-            "sub_answers_graph_initial",
-            BaseQAState(
-                sub_question_str=initial_sub_question["sub_question_str"],
-                sub_question_search_queries=initial_sub_question[
-                    "sub_question_search_queries"
-                ],
-                sub_question_nr=initial_sub_question["sub_question_nr"],
-                primary_llm=state["primary_llm"],
-                fast_llm=state["fast_llm"],
-                graph_start_time=state["graph_start_time"],
-            ),
-        )
-        for initial_sub_question in state["initial_sub_questions"]
-    ]
-
-
-def continue_to_answer_sub_questions(state: QAState) -> Union[Hashable, list[Hashable]]:
-    # Routes re-written queries to the (parallel) retrieval steps
-    # Notice the 'Send()' API that takes care of the parallelization
-    return [
-        Send(
-            "sub_answers_graph",
-            ResearchQAState(
-                sub_question=sub_question["sub_question_str"],
-                sub_question_nr=sub_question["sub_question_nr"],
-                graph_start_time=state["graph_start_time"],
-                primary_llm=state["primary_llm"],
-                fast_llm=state["fast_llm"],
-            ),
-        )
-        for sub_question in state["sub_questions"]
-    ]
-
-
-def continue_to_deep_answer(state: QAState) -> Union[Hashable, list[Hashable]]:
-    print("---GO TO DEEP ANSWER OR END---")
-
-    base_answer = state["base_answer"]
-
-    question = state["original_question"]
-
-    BASE_CHECK_MESSAGE = [
-        HumanMessage(
-            content=BASE_CHECK_PROMPT.format(question=question, base_answer=base_answer)
-        )
-    ]
-
-    model = state["fast_llm"]
-    response = model.invoke(BASE_CHECK_MESSAGE)
-
-    print(f"CAN WE CONTINUE W/O GENERATING A DEEP ANSWER? - {response.pretty_repr()}")
-
-    if response.pretty_repr() == "no":
-        return "decompose"
-    else:
-        return "end"

backend/danswer/agent_search/primary_graph/graph_builder.py

@@ -1,171 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from danswer.agent_search.core_qa_graph.graph_builder import build_core_qa_graph
-from danswer.agent_search.deep_qa_graph.graph_builder import build_deep_qa_graph
-from danswer.agent_search.primary_graph.edges import continue_to_answer_sub_questions
-from danswer.agent_search.primary_graph.edges import continue_to_deep_answer
-from danswer.agent_search.primary_graph.edges import continue_to_initial_sub_questions
-from danswer.agent_search.primary_graph.nodes.base_wait import base_wait
-from danswer.agent_search.primary_graph.nodes.combine_retrieved_docs import (
-    combine_retrieved_docs,
-)
-from danswer.agent_search.primary_graph.nodes.custom_retrieve import custom_retrieve
-from danswer.agent_search.primary_graph.nodes.decompose import decompose
-from danswer.agent_search.primary_graph.nodes.deep_answer_generation import (
-    deep_answer_generation,
-)
-from danswer.agent_search.primary_graph.nodes.dummy_start import dummy_start
-from danswer.agent_search.primary_graph.nodes.entity_term_extraction import (
-    entity_term_extraction,
-)
-from danswer.agent_search.primary_graph.nodes.final_stuff import final_stuff
-from danswer.agent_search.primary_graph.nodes.generate_initial import generate_initial
-from danswer.agent_search.primary_graph.nodes.main_decomp_base import main_decomp_base
-from danswer.agent_search.primary_graph.nodes.rewrite import rewrite
-from danswer.agent_search.primary_graph.nodes.sub_qa_level_aggregator import (
-    sub_qa_level_aggregator,
-)
-from danswer.agent_search.primary_graph.nodes.sub_qa_manager import sub_qa_manager
-from danswer.agent_search.primary_graph.nodes.verifier import verifier
-from danswer.agent_search.primary_graph.states import QAState
-
-
-def build_core_graph() -> StateGraph:
-    # Define the nodes we will cycle between
-    core_answer_graph = StateGraph(state_schema=QAState)
-
-    ### Add Nodes ###
-    core_answer_graph.add_node(node="dummy_start", 
-                               action=dummy_start)
-
-    # Re-writing the question
-    core_answer_graph.add_node(node="rewrite", 
-                               action=rewrite)
-
-    # The retrieval step
-    core_answer_graph.add_node(node="custom_retrieve", 
-                               action=custom_retrieve)
-
-    # Combine and dedupe retrieved docs.
-    core_answer_graph.add_node(
-        node="combine_retrieved_docs", 
-        action=combine_retrieved_docs
-    )
-
-    # Extract entities, terms and relationships
-    core_answer_graph.add_node(
-        node="entity_term_extraction", 
-        action=entity_term_extraction
-    )
-
-    # Verifying that a retrieved doc is relevant
-    core_answer_graph.add_node(node="verifier", 
-                               action=verifier)
-
-    # Initial question decomposition
-    core_answer_graph.add_node(node="main_decomp_base", 
-                               action=main_decomp_base)
-
-    # Build the base QA sub-graph and compile it
-    compiled_core_qa_graph = build_core_qa_graph().compile()
-    # Add the compiled base QA sub-graph as a node to the core graph
-    core_answer_graph.add_node(
-        node="sub_answers_graph_initial", 
-        action=compiled_core_qa_graph
-    )
-
-    # Checking whether the initial answer is in the ballpark
-    core_answer_graph.add_node(node="base_wait", 
-                               action=base_wait)
-
-    # Decompose the question into sub-questions
-    core_answer_graph.add_node(node="decompose", 
-                               action=decompose)
-
-    # Manage the sub-questions
-    core_answer_graph.add_node(node="sub_qa_manager", 
-                               action=sub_qa_manager)
-
-    # Build the research QA sub-graph and compile it
-    compiled_deep_qa_graph = build_deep_qa_graph().compile()
-    # Add the compiled research QA sub-graph as a node to the core graph
-    core_answer_graph.add_node(node="sub_answers_graph", 
-                               action=compiled_deep_qa_graph)
-
-    # Aggregate the sub-questions
-    core_answer_graph.add_node(
-        node="sub_qa_level_aggregator", 
-        action=sub_qa_level_aggregator
-    )
-
-    # aggregate sub questions and answers
-    core_answer_graph.add_node(
-        node="deep_answer_generation", 
-        action=deep_answer_generation
-    )
-
-    # A final clean-up step
-    core_answer_graph.add_node(node="final_stuff", 
-                               action=final_stuff)
-
-    # Generating a response after we know the documents are relevant
-    core_answer_graph.add_node(node="generate_initial", 
-                               action=generate_initial)
-
-    ### Add Edges ###
-
-    # start the initial sub-question decomposition
-    core_answer_graph.add_edge(start_key=START, 
-                               end_key="main_decomp_base")
-
-    core_answer_graph.add_conditional_edges(
-        source="main_decomp_base",
-        path=continue_to_initial_sub_questions,
-    )
-
-    # use the retrieved information to generate the answer
-    core_answer_graph.add_edge(
-        start_key=["verifier", "sub_answers_graph_initial"], 
-        end_key="generate_initial"
-    )
-    core_answer_graph.add_edge(start_key="generate_initial", 
-                               end_key="base_wait")
-
-    core_answer_graph.add_conditional_edges(
-        source="base_wait",
-        path=continue_to_deep_answer,
-        path_map={"decompose": "entity_term_extraction", "end": "final_stuff"},
-    )
-
-    core_answer_graph.add_edge(start_key="entity_term_extraction", end_key="decompose")
-
-    core_answer_graph.add_edge(start_key="decompose", 
-                               end_key="sub_qa_manager")
-    core_answer_graph.add_conditional_edges(
-        source="sub_qa_manager",
-        path=continue_to_answer_sub_questions,
-    )
-
-    core_answer_graph.add_edge(
-        start_key="sub_answers_graph",
-          end_key="sub_qa_level_aggregator"
-    )
-
-    core_answer_graph.add_edge(
-        start_key="sub_qa_level_aggregator", 
-        end_key="deep_answer_generation"
-    )
-
-    core_answer_graph.add_edge(
-        start_key="deep_answer_generation", 
-        end_key="final_stuff"
-    )
-
-    core_answer_graph.add_edge(start_key="final_stuff", 
-                               end_key=END)
-    
-    core_answer_graph.compile()
-
-    return core_answer_graph

backend/danswer/agent_search/primary_graph/nodes/base_wait.py

@@ -1,27 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def base_wait(state: QAState) -> dict[str, Any]:
-    """
-    Ensures that all required steps are completed before proceeding to the next step
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: {} (no operation, just logging)
-    """
-
-    print("---Base Wait ---")
-    node_start_time = datetime.now()
-    return {
-        "log_messages": generate_log_message(
-            message="core - base_wait",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/combine_retrieved_docs.py

@@ -1,36 +0,0 @@
-from collections.abc import Sequence
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.context.search.models import InferenceSection
-
-
-def combine_retrieved_docs(state: QAState) -> dict[str, Any]:
-    """
-    Dedupe the retrieved docs.
-    """
-    node_start_time = datetime.now()
-
-    base_retrieval_docs: Sequence[InferenceSection] = state["base_retrieval_docs"]
-
-    print(f"Number of docs from steps: {len(base_retrieval_docs)}")
-    dedupe_docs: list[InferenceSection] = []
-    for base_retrieval_doc in base_retrieval_docs:
-        if not any(
-            base_retrieval_doc.center_chunk.document_id == doc.center_chunk.document_id
-            for doc in dedupe_docs
-        ):
-            dedupe_docs.append(base_retrieval_doc)
-
-    print(f"Number of deduped docs: {len(dedupe_docs)}")
-
-    return {
-        "deduped_retrieval_docs": dedupe_docs,
-        "log_messages": generate_log_message(
-            message="core - combine_retrieved_docs (dedupe)",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/custom_retrieve.py

@@ -1,52 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import RetrieverState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.context.search.models import InferenceSection
-from danswer.context.search.models import SearchRequest
-from danswer.context.search.pipeline import SearchPipeline
-from danswer.db.engine import get_session_context_manager
-from danswer.llm.factory import get_default_llms
-
-
-def custom_retrieve(state: RetrieverState) -> dict[str, Any]:
-    """
-    Retrieve documents
-
-    Args:
-        retriever_state (dict): The current graph state
-
-    Returns:
-        state (dict): New key added to state, documents, that contains retrieved documents
-    """
-    print("---RETRIEVE---")
-
-    node_start_time = datetime.now()
-
-    query = state["rewritten_query"]
-
-    # Retrieval
-    # TODO: add the actual retrieval, probably from search_tool.run()
-    llm, fast_llm = get_default_llms()
-    with get_session_context_manager() as db_session:
-        top_sections = SearchPipeline(
-            search_request=SearchRequest(
-                query=query,
-            ),
-            user=None,
-            llm=llm,
-            fast_llm=fast_llm,
-            db_session=db_session,
-        ).reranked_sections
-        print(len(top_sections))
-    documents: list[InferenceSection] = []
-
-    return {
-        "base_retrieval_docs": documents,
-        "log_messages": generate_log_message(
-            message="core - custom_retrieve",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/decompose.py

@@ -1,78 +0,0 @@
-import json
-import re
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.prompts import DEEP_DECOMPOSE_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_entity_term_extraction
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def decompose(state: QAState) -> dict[str, Any]:
-    """ """
-
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-    base_answer = state["base_answer"]
-
-    # get the entity term extraction dict and properly format it
-    entity_term_extraction_dict = state["retrieved_entities_relationships"][
-        "retrieved_entities_relationships"
-    ]
-
-    entity_term_extraction_str = format_entity_term_extraction(
-        entity_term_extraction_dict
-    )
-
-    initial_question_answers = state["initial_sub_qas"]
-
-    addressed_question_list = [
-        x["sub_question"]
-        for x in initial_question_answers
-        if x["sub_answer_check"] == "yes"
-    ]
-    failed_question_list = [
-        x["sub_question"]
-        for x in initial_question_answers
-        if x["sub_answer_check"] == "no"
-    ]
-
-    msg = [
-        HumanMessage(
-            content=DEEP_DECOMPOSE_PROMPT.format(
-                question=question,
-                entity_term_extraction_str=entity_term_extraction_str,
-                base_answer=base_answer,
-                answered_sub_questions="\n - ".join(addressed_question_list),
-                failed_sub_questions="\n - ".join(failed_question_list),
-            ),
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    cleaned_response = re.sub(r"```json\n|\n```", "", response.pretty_repr())
-    parsed_response = json.loads(cleaned_response)
-
-    sub_questions_dict = {}
-    for sub_question_nr, sub_question_dict in enumerate(
-        parsed_response["sub_questions"]
-    ):
-        sub_question_dict["answered"] = False
-        sub_question_dict["verified"] = False
-        sub_questions_dict[sub_question_nr] = sub_question_dict
-
-    return {
-        "decomposed_sub_questions_dict": sub_questions_dict,
-        "log_messages": generate_log_message(
-            message="deep - decompose",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/deep_answer_generation.py

@@ -1,61 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.prompts import COMBINED_CONTEXT
-from danswer.agent_search.shared_graph_utils.prompts import MODIFIED_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.agent_search.shared_graph_utils.utils import normalize_whitespace
-
-
-# aggregate sub questions and answers
-def deep_answer_generation(state: QAState) -> dict[str, Any]:
-    """
-    Generate answer
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-         dict: The updated state with re-phrased question
-    """
-    print("---DEEP GENERATE---")
-
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-    docs = state["deduped_retrieval_docs"]
-
-    deep_answer_context = state["core_answer_dynamic_context"]
-
-    print(f"Number of verified retrieval docs - deep: {len(docs)}")
-
-    combined_context = normalize_whitespace(
-        COMBINED_CONTEXT.format(
-            deep_answer_context=deep_answer_context, formated_docs=format_docs(docs)
-        )
-    )
-
-    msg = [
-        HumanMessage(
-            content=MODIFIED_RAG_PROMPT.format(
-                question=question, combined_context=combined_context
-            )
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    return {
-        "deep_answer": response.content,
-        "log_messages": generate_log_message(
-            message="deep - deep answer generation",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/dummy_start.py

@@ -1,11 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import QAState
-
-
-def dummy_start(state: QAState) -> dict[str, Any]:
-    """
-    Dummy node to set the start time
-    """
-    return {"start_time": datetime.now()}

backend/danswer/agent_search/primary_graph/nodes/entity_term_extraction.py

@@ -1,51 +0,0 @@
-import json
-import re
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-from langchain_core.messages import merge_message_runs
-
-from danswer.agent_search.primary_graph.prompts import ENTITY_TERM_PROMPT
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-from danswer.llm.factory import get_default_llms
-
-
-def entity_term_extraction(state: QAState) -> dict[str, Any]:
-    """Extract entities and terms from the question and context"""
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-    docs = state["deduped_retrieval_docs"]
-
-    doc_context = format_docs(docs)
-
-    msg = [
-        HumanMessage(
-            content=ENTITY_TERM_PROMPT.format(question=question, context=doc_context),
-        )
-    ]
-    _, fast_llm = get_default_llms()
-    # Grader
-    llm_response_list = list(
-        fast_llm.stream(
-            prompt=msg,
-            # structured_response_format={"type": "json_object", "schema": RewrittenQueries.model_json_schema()},
-            # structured_response_format=RewrittenQueries.model_json_schema(),
-        )
-    )
-    llm_response = merge_message_runs(llm_response_list, chunk_separator="")[0].content
-
-    cleaned_response = re.sub(r"```json\n|\n```", "", llm_response)
-    parsed_response = json.loads(cleaned_response)
-
-    return {
-        "retrieved_entities_relationships": parsed_response,
-        "log_messages": generate_log_message(
-            message="deep - entity term extraction",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/final_stuff.py

@@ -1,85 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def final_stuff(state: QAState) -> dict[str, Any]:
-    """
-    Invokes the agent model to generate a response based on the current state. Given
-    the question, it will decide to retrieve using the retriever tool, or simply end.
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: The updated state with the agent response appended to messages
-    """
-    print("---FINAL---")
-    node_start_time = datetime.now()
-
-    messages = state["log_messages"]
-    time_ordered_messages = [x.pretty_repr() for x in messages]
-    time_ordered_messages.sort()
-
-    print("Message Log:")
-    print("\n".join(time_ordered_messages))
-
-    initial_sub_qas = state["initial_sub_qas"]
-    initial_sub_qa_list = []
-    for initial_sub_qa in initial_sub_qas:
-        if initial_sub_qa["sub_answer_check"] == "yes":
-            initial_sub_qa_list.append(
-                f'  Question:\n  {initial_sub_qa["sub_question"]}\n  --\n  Answer:\n  {initial_sub_qa["sub_answer"]}\n  -----'
-            )
-
-    initial_sub_qa_context = "\n".join(initial_sub_qa_list)
-
-    log_message = generate_log_message(
-        message="all - final_stuff",
-        node_start_time=node_start_time,
-        graph_start_time=state["graph_start_time"],
-    )
-
-    print(log_message)
-    print("--------------------------------")
-
-    base_answer = state["base_answer"]
-
-    print(f"Final Base Answer:\n{base_answer}")
-    print("--------------------------------")
-    print(f"Initial Answered Sub Questions:\n{initial_sub_qa_context}")
-    print("--------------------------------")
-
-    if not state.get("deep_answer"):
-        print("No Deep Answer was required")
-        return {
-            "log_messages": log_message,
-        }
-
-    deep_answer = state["deep_answer"]
-    sub_qas = state["sub_qas"]
-    sub_qa_list = []
-    for sub_qa in sub_qas:
-        if sub_qa["sub_answer_check"] == "yes":
-            sub_qa_list.append(
-                f'  Question:\n  {sub_qa["sub_question"]}\n  --\n  Answer:\n  {sub_qa["sub_answer"]}\n  -----'
-            )
-
-    sub_qa_context = "\n".join(sub_qa_list)
-
-    print(f"Final Base Answer:\n{base_answer}")
-    print("--------------------------------")
-    print(f"Final Deep Answer:\n{deep_answer}")
-    print("--------------------------------")
-    print("Sub Questions and Answers:")
-    print(sub_qa_context)
-
-    return {
-        "log_messages": generate_log_message(
-            message="all - final_stuff",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/generate.py

@@ -1,52 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.prompts import BASE_RAG_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def generate(state: QAState) -> dict[str, Any]:
-    """
-    Generate answer
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-         dict: The updated state with re-phrased question
-    """
-    print("---GENERATE---")
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-    docs = state["deduped_retrieval_docs"]
-
-    print(f"Number of verified retrieval docs: {len(docs)}")
-
-    msg = [
-        HumanMessage(
-            content=BASE_RAG_PROMPT.format(question=question, context=format_docs(docs))
-        )
-    ]
-
-    # Grader
-    llm = state["fast_llm"]
-    response = list(
-        llm.stream(
-            prompt=msg,
-            structured_response_format=None,
-        )
-    )
-
-    return {
-        "base_answer": response[0].pretty_repr(),
-        "log_messages": generate_log_message(
-            message="core - generate",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/generate_initial.py

@@ -1,72 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.prompts import INITIAL_RAG_PROMPT
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import format_docs
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def generate_initial(state: QAState) -> dict[str, Any]:
-    """
-    Generate answer
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-         dict: The updated state with re-phrased question
-    """
-    print("---GENERATE INITIAL---")
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-    docs = state["deduped_retrieval_docs"]
-    print(f"Number of verified retrieval docs - base: {len(docs)}")
-
-    sub_question_answers = state["initial_sub_qas"]
-
-    sub_question_answers_list = []
-
-    _SUB_QUESTION_ANSWER_TEMPLATE = """
-    Sub-Question:\n  - {sub_question}\n  --\nAnswer:\n  - {sub_answer}\n\n
-    """
-    for sub_question_answer_dict in sub_question_answers:
-        if (
-            sub_question_answer_dict["sub_answer_check"] == "yes"
-            and len(sub_question_answer_dict["sub_answer"]) > 0
-            and sub_question_answer_dict["sub_answer"] != "I don't know"
-        ):
-            sub_question_answers_list.append(
-                _SUB_QUESTION_ANSWER_TEMPLATE.format(
-                    sub_question=sub_question_answer_dict["sub_question"],
-                    sub_answer=sub_question_answer_dict["sub_answer"],
-                )
-            )
-
-    sub_question_answer_str = "\n\n------\n\n".join(sub_question_answers_list)
-
-    msg = [
-        HumanMessage(
-            content=INITIAL_RAG_PROMPT.format(
-                question=question,
-                context=format_docs(docs),
-                answered_sub_questions=sub_question_answer_str,
-            )
-        )
-    ]
-
-    # Grader
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    return {
-        "base_answer": response.pretty_repr(),
-        "log_messages": generate_log_message(
-            message="core - generate initial",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/main_decomp_base.py

@@ -1,64 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.prompts import INITIAL_DECOMPOSITION_PROMPT
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import clean_and_parse_list_string
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def main_decomp_base(state: QAState) -> dict[str, Any]:
-    """
-    Perform an initial question decomposition, incl. one search term
-
-    Args:
-        state (messages): The current state
-
-    Returns:
-        dict: The updated state with initial decomposition
-    """
-
-    print("---INITIAL DECOMP---")
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-
-    msg = [
-        HumanMessage(
-            content=INITIAL_DECOMPOSITION_PROMPT.format(question=question),
-        )
-    ]
-
-    # Get the rewritten queries in a defined format
-    model = state["fast_llm"]
-    response = model.invoke(msg)
-
-    content = response.pretty_repr()
-    list_of_subquestions = clean_and_parse_list_string(content)
-
-    decomp_list = []
-
-    for sub_question_nr, sub_question in enumerate(list_of_subquestions):
-        sub_question_str = sub_question["sub_question"].strip()
-        # temporarily
-        sub_question_search_queries = [sub_question["search_term"]]
-
-        decomp_list.append(
-            {
-                "sub_question_str": sub_question_str,
-                "sub_question_search_queries": sub_question_search_queries,
-                "sub_question_nr": sub_question_nr,
-            }
-        )
-
-    return {
-        "initial_sub_questions": decomp_list,
-        "sub_query_start_time": node_start_time,
-        "log_messages": generate_log_message(
-            message="core - initial decomp",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/rewrite.py

@@ -1,55 +0,0 @@
-import json
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.models import RewrittenQueries
-from danswer.agent_search.shared_graph_utils.prompts import REWRITE_PROMPT_MULTI
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def rewrite(state: QAState) -> dict[str, Any]:
-    """
-    Transform the initial question into more suitable search queries.
-
-    Args:
-        qa_state (messages): The current state
-
-    Returns:
-        dict: The updated state with re-phrased question
-    """
-    print("---STARTING GRAPH---")
-    graph_start_time = datetime.now()
-
-    print("---TRANSFORM QUERY---")
-    node_start_time = datetime.now()
-
-    question = state["original_question"]
-
-    msg = [
-        HumanMessage(
-            content=REWRITE_PROMPT_MULTI.format(question=question),
-        )
-    ]
-
-    # Get the rewritten queries in a defined format
-    fast_llm = state["fast_llm"]
-    llm_response = list(
-        fast_llm.stream(
-            prompt=msg,
-            structured_response_format=RewrittenQueries.model_json_schema(),
-        )
-    )
-
-    formatted_response: RewrittenQueries = json.loads(llm_response[0].pretty_repr())
-
-    return {
-        "rewritten_queries": formatted_response.rewritten_queries,
-        "log_messages": generate_log_message(
-            message="core - rewrite",
-            node_start_time=node_start_time,
-            graph_start_time=graph_start_time,
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/sub_qa_level_aggregator.py

@@ -1,39 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-# aggregate sub questions and answers
-def sub_qa_level_aggregator(state: QAState) -> dict[str, Any]:
-    sub_qas = state["sub_qas"]
-
-    node_start_time = datetime.now()
-
-    dynamic_context_list = [
-        "Below you will find useful information to answer the original question:"
-    ]
-    checked_sub_qas = []
-
-    for core_answer_sub_qa in sub_qas:
-        question = core_answer_sub_qa["sub_question"]
-        answer = core_answer_sub_qa["sub_answer"]
-        verified = core_answer_sub_qa["sub_answer_check"]
-
-        if verified == "yes":
-            dynamic_context_list.append(
-                f"Question:\n{question}\n\nAnswer:\n{answer}\n\n---\n\n"
-            )
-            checked_sub_qas.append({"sub_question": question, "sub_answer": answer})
-    dynamic_context = "\n".join(dynamic_context_list)
-
-    return {
-        "core_answer_dynamic_context": dynamic_context,
-        "checked_sub_qas": checked_sub_qas,
-        "log_messages": generate_log_message(
-            message="deep - sub qa level aggregator",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/sub_qa_manager.py

@@ -1,28 +0,0 @@
-from datetime import datetime
-from typing import Any
-
-from danswer.agent_search.primary_graph.states import QAState
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def sub_qa_manager(state: QAState) -> dict[str, Any]:
-    """ """
-
-    node_start_time = datetime.now()
-
-    sub_questions_dict = state["decomposed_sub_questions_dict"]
-
-    sub_questions = {}
-
-    for sub_question_nr, sub_question_dict in sub_questions_dict.items():
-        sub_questions[sub_question_nr] = sub_question_dict["sub_question"]
-
-    return {
-        "sub_questions": sub_questions,
-        "num_new_question_iterations": 0,
-        "log_messages": generate_log_message(
-            message="deep - sub qa manager",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/nodes/verifier.py

@@ -1,59 +0,0 @@
-import json
-from datetime import datetime
-from typing import Any
-
-from langchain_core.messages import HumanMessage
-
-from danswer.agent_search.primary_graph.states import VerifierState
-from danswer.agent_search.shared_graph_utils.models import BinaryDecision
-from danswer.agent_search.shared_graph_utils.prompts import VERIFIER_PROMPT
-from danswer.agent_search.shared_graph_utils.utils import generate_log_message
-
-
-def verifier(state: VerifierState) -> dict[str, Any]:
-    """
-    Check whether the document is relevant for the original user question
-
-    Args:
-        state (VerifierState): The current state
-
-    Returns:
-        dict: ict: The updated state with the final decision
-    """
-
-    print("---VERIFY QUTPUT---")
-    node_start_time = datetime.now()
-
-    question = state["question"]
-    document_content = state["document"].combined_content
-
-    msg = [
-        HumanMessage(
-            content=VERIFIER_PROMPT.format(
-                question=question, document_content=document_content
-            )
-        )
-    ]
-
-    # Grader
-    llm = state["fast_llm"]
-    response = list(
-        llm.stream(
-            prompt=msg,
-            structured_response_format=BinaryDecision.model_json_schema(),
-        )
-    )
-
-    raw_response = json.loads(response[0].pretty_repr())
-    formatted_response = BinaryDecision.model_validate(raw_response)
-
-    return {
-        "deduped_retrieval_docs": [state["document"]]
-        if formatted_response.decision == "yes"
-        else [],
-        "log_messages": generate_log_message(
-            message=f"core - verifier: {formatted_response.decision}",
-            node_start_time=node_start_time,
-            graph_start_time=state["graph_start_time"],
-        ),
-    }

backend/danswer/agent_search/primary_graph/prompts.py

@@ -1,86 +0,0 @@
-INITIAL_DECOMPOSITION_PROMPT = """ \n
-    Please decompose an initial user question into not more than 4 appropriate sub-questions that help to
-    answer the original question. The purpose for this decomposition is to isolate individulal entities
-    (i.e., 'compare sales of company A and company B' -> 'what are sales for company A' + 'what are sales
-    for company B'), split ambiguous terms (i.e., 'what is our success with company A' -> 'what are our
-    sales with company A' + 'what is our market share with company A' + 'is company A a reference customer
-    for us'), etc. Each sub-question should be realistically be answerable by a good RAG system. \n
-
-    For each sub-question, please also create one search term that can be used to retrieve relevant
-    documents from a document store.
-
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Please formulate your answer as a list of json objects with the following format:
-
-   [{{"sub_question": <sub-question>, "search_term": <search term>}}, ...]
-
-    Answer:
-    """
-
-INITIAL_RAG_PROMPT = """ \n
-    You are an assistant for question-answering tasks. Use the information provided below - and only the
-    provided information - to answer the provided question.
-
-    The information provided below consists of:
-     1) a number of answered sub-questions - these are very important(!) and definitely should be
-     considered to answer the question.
-     2) a number of documents that were also deemed relevant for the question.
-
-    If you don't know the answer or if the provided information is empty or insufficient, just say
-    "I don't know". Do not use your internal knowledge!
-
-    Again, only use the provided informationand do not use your internal knowledge! It is a matter of life
-    and death that you do NOT use your internal knowledge, just the provided information!
-
-    Try to keep your answer concise.
-
-    And here is the question and the provided information:
-    \n
-    \nQuestion:\n {question}
-
-    \nAnswered Sub-questions:\n {answered_sub_questions}
-
-    \nContext:\n {context} \n\n
-    \n\n
-
-    Answer:"""
-
-ENTITY_TERM_PROMPT = """ \n
-    Based on the original question and the context retieved from a dataset, please generate a list of
-    entities (e.g. companies, organizations, industries, products, locations, etc.), terms and concepts
-    (e.g. sales, revenue, etc.) that are relevant for the question, plus their relations to each other.
-
-    \n\n
-    Here is the original question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-   And here is the context retrieved:
-    \n ------- \n
-    {context}
-    \n ------- \n
-
-    Please format your answer as a json object in the following format:
-
-    {{"retrieved_entities_relationships": {{
-        "entities": [{{
-            "entity_name": <assign a name for the entity>,
-            "entity_type": <specify a short type name for the entity, such as 'company', 'location',...>
-        }}],
-        "relationships": [{{
-            "name": <assign a name for the relationship>,
-            "type": <specify a short type name for the relationship, such as 'sales_to', 'is_location_of',...>,
-            "entities": [<related entity name 1>, <related entity name 2>]
-        }}],
-        "terms": [{{
-            "term_name": <assign a name for the term>,
-            "term_type": <specify a short type name for the term, such as 'revenue', 'market_share',...>,
-            "similar_to": <list terms that are similar to this term>
-        }}]
-    }}
-    }}
-   """

backend/danswer/agent_search/primary_graph/states.py

@@ -1,73 +0,0 @@
-import operator
-from collections.abc import Sequence
-from datetime import datetime
-from typing import Annotated
-from typing import TypedDict
-
-from langchain_core.messages import BaseMessage
-from langgraph.graph.message import add_messages
-
-from danswer.agent_search.shared_graph_utils.models import RewrittenQueries
-from danswer.context.search.models import InferenceSection
-
-
-class QAState(TypedDict):
-    # The 'main' state of the answer graph
-    original_question: str
-    graph_start_time: datetime
-    # start time for parallel initial sub-questionn thread
-    sub_query_start_time: datetime
-    log_messages: Annotated[Sequence[BaseMessage], add_messages]
-    rewritten_queries: RewrittenQueries
-    sub_questions: list[dict]
-    initial_sub_questions: list[dict]
-    ranked_subquestion_ids: list[int]
-    decomposed_sub_questions_dict: dict
-    rejected_sub_questions: Annotated[list[str], operator.add]
-    rejected_sub_questions_handled: bool
-    sub_qas: Annotated[Sequence[dict], operator.add]
-    initial_sub_qas: Annotated[Sequence[dict], operator.add]
-    checked_sub_qas: Annotated[Sequence[dict], operator.add]
-    base_retrieval_docs: Annotated[Sequence[InferenceSection], operator.add]
-    deduped_retrieval_docs: Annotated[Sequence[InferenceSection], operator.add]
-    reranked_retrieval_docs: Annotated[Sequence[InferenceSection], operator.add]
-    retrieved_entities_relationships: dict
-    questions_context: list[dict]
-    qa_level: int
-    top_chunks: list[InferenceSection]
-    sub_question_top_chunks: Annotated[Sequence[dict], operator.add]
-    num_new_question_iterations: int
-    core_answer_dynamic_context: str
-    dynamic_context: str
-    initial_base_answer: str
-    base_answer: str
-    deep_answer: str
-
-
-class QAOuputState(TypedDict):
-    # The 'main' output state of the answer graph. Removes all the intermediate states
-    original_question: str
-    log_messages: Annotated[Sequence[BaseMessage], add_messages]
-    sub_questions: list[dict]
-    sub_qas: Annotated[Sequence[dict], operator.add]
-    initial_sub_qas: Annotated[Sequence[dict], operator.add]
-    checked_sub_qas: Annotated[Sequence[dict], operator.add]
-    reranked_retrieval_docs: Annotated[Sequence[InferenceSection], operator.add]
-    retrieved_entities_relationships: dict
-    top_chunks: list[InferenceSection]
-    sub_question_top_chunks: Annotated[Sequence[dict], operator.add]
-    base_answer: str
-    deep_answer: str
-
-
-class RetrieverState(TypedDict):
-    # The state for the parallel Retrievers. They each need to see only one query
-    rewritten_query: str
-    graph_start_time: datetime
-
-
-class VerifierState(TypedDict):
-    # The state for the parallel verification step.  Each node execution need to see only one question/doc pair
-    document: InferenceSection
-    question: str
-    graph_start_time: datetime

backend/danswer/agent_search/run_graph.py

@@ -1,22 +0,0 @@
-from danswer.agent_search.primary_graph.graph_builder import build_core_graph
-from danswer.llm.answering.answer import AnswerStream
-from danswer.llm.interfaces import LLM
-from danswer.tools.tool import Tool
-
-
-def run_graph(
-    query: str,
-    llm: LLM,
-    tools: list[Tool],
-) -> AnswerStream:
-    graph = build_core_graph()
-
-    inputs = {
-        "original_question": query,
-        "messages": [],
-        "tools": tools,
-        "llm": llm,
-    }
-    compiled_graph = graph.compile()
-    output = compiled_graph.invoke(input=inputs)
-    yield from output

backend/danswer/agent_search/shared_graph_utils/models.py

@@ -1,16 +0,0 @@
-from typing import Literal
-
-from pydantic import BaseModel
-
-
-# Pydantic models for structured outputs
-class RewrittenQueries(BaseModel):
-    rewritten_queries: list[str]
-
-
-class BinaryDecision(BaseModel):
-    decision: Literal["yes", "no"]
-
-
-class SubQuestions(BaseModel):
-    sub_questions: list[str]

backend/danswer/agent_search/shared_graph_utils/prompts.py

@@ -1,342 +0,0 @@
-REWRITE_PROMPT_MULTI_ORIGINAL = """ \n
-    Please convert an initial user question into a 2-3 more appropriate short and pointed search queries for retrievel from a
-    document store. Particularly, try to think about resolving ambiguities and make the search queries more specific,
-    enabling the system to search more broadly.
-    Also, try to make the search queries not redundant, i.e. not too similar! \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Formulate the queries separated by '--' (Do not say 'Query 1: ...', just write the querytext): """
-
-
-REWRITE_PROMPT_MULTI = """ \n
-    Please create a list of 2-3 sample documents that could answer an original question. Each document
-    should be about as long as the original question. \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Formulate the sample documents separated by '--' (Do not say 'Document 1: ...', just write the text): """
-
-BASE_RAG_PROMPT = """ \n
-    You are an assistant for question-answering tasks. Use the context provided below - and only the
-    provided context - to answer the question. If you don't know the answer or if the provided context is
-    empty, just say "I don't know". Do not use your internal knowledge!
-
-    Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
-    question based on the context, say "I don't know". It is a matter of life and death that you do NOT
-    use your internal knowledge, just the provided information!
-
-    Use three sentences maximum and keep the answer concise.
-    answer concise.\nQuestion:\n {question} \nContext:\n {context} \n\n
-    \n\n
-    Answer:"""
-
-BASE_CHECK_PROMPT = """ \n
-    Please check whether 1) the suggested answer seems to fully address the original question AND 2)the
-    original question requests a simple, factual answer, and there are no ambiguities, judgements,
-    aggregations, or any other complications that may require extra context. (I.e., if the question is
-    somewhat addressed, but the answer would benefit from more context, then answer with 'no'.)
-
-    Please only answer with 'yes' or 'no' \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Here is the proposed answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-    Please answer with yes or no:"""
-
-VERIFIER_PROMPT = """ \n
-    Please check whether the document seems to be relevant for the answer of the original question. Please
-    only answer with 'yes' or 'no' \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-    Here is the document text:
-    \n ------- \n
-    {document_content}
-    \n ------- \n
-    Please answer with yes or no:"""
-
-INITIAL_DECOMPOSITION_PROMPT_BASIC = """ \n
-    Please decompose an initial user question into not more than 4 appropriate sub-questions that help to
-    answer the original question. The purpose for this decomposition is to isolate individulal entities
-    (i.e., 'compare sales of company A and company B' -> 'what are sales for company A' + 'what are sales
-    for company B'), split ambiguous terms (i.e., 'what is our success with company A' -> 'what are our
-    sales with company A' + 'what is our market share with company A' + 'is company A a reference customer
-    for us'), etc. Each sub-question should be realistically be answerable by a good RAG system. \n
-
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Please formulate your answer as a list of subquestions:
-
-    Answer:
-    """
-
-REWRITE_PROMPT_SINGLE = """ \n
-    Please convert an initial user question into a more appropriate search query for retrievel from a
-    document store. \n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Formulate the query: """
-
-MODIFIED_RAG_PROMPT = """You are an assistant for question-answering tasks. Use the context provided below
-    - and only this context - to answer the question. If you don't know the answer, just say "I don't know".
-    Use three sentences maximum and keep the answer concise.
-    Pay also particular attention to the sub-questions and their answers, at least it may enrich the answer.
-    Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
-    question based on the context, say "I don't know". It is a matter of life and death that you do NOT
-    use your internal knowledge, just the provided information!
-
-    \nQuestion: {question}
-    \nContext: {combined_context} \n
-
-    Answer:"""
-
-ORIG_DEEP_DECOMPOSE_PROMPT = """ \n
-    An initial user question needs to be answered. An initial answer has been provided but it wasn't quite
-    good enough. Also, some sub-questions had been answered and this information has been used to provide
-    the initial answer. Some other subquestions may have been suggested based on little knowledge, but they
-    were not directly answerable. Also, some entities, relationships and terms are givenm to you so that
-    you have an idea of how the avaiolable data looks like.
-
-    Your role is to generate 3-5 new sub-questions that would help to answer the initial question,
-    considering:
-
-    1) The initial question
-    2) The initial answer that was found to be unsatisfactory
-    3) The sub-questions that were answered
-    4) The sub-questions that were suggested but not answered
-    5) The entities, relationships and terms that were extracted from the context
-
-    The individual questions should be answerable by a good RAG system.
-    So a good idea would be to use the sub-questions to resolve ambiguities and/or to separate the
-    question for different entities that may be involved in the original question, but in a way that does
-    not duplicate questions that were already tried.
-
-    Additional Guidelines:
-    - The sub-questions should be specific to the question and provide richer context for the question,
-    resolve ambiguities, or address shortcoming of the initial answer
-    - Each sub-question - when answered - should be relevant for the answer to the original question
-    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
-    other complications that may require extra context.
-    - The sub-questions MUST have the full context of the original question so that it can be executed by
-    a RAG system independently without the original question available
-      (Example:
-        - initial question: "What is the capital of France?"
-        - bad sub-question: "What is the name of the river there?"
-        - good sub-question: "What is the name of the river that flows through Paris?"
-    - For each sub-question, please provide a short explanation for why it is a good sub-question. So
-    generate a list of dictionaries with the following format:
-      [{{"sub_question": <sub-question>, "explanation": <explanation>, "search_term": <rewrite the
-      sub-question using as a search phrase for the document store>}}, ...]
-
-    \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Here is the initial sub-optimal answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-
-    Here are the sub-questions that were answered:
-    \n ------- \n
-    {answered_sub_questions}
-    \n ------- \n
-
-    Here are the sub-questions that were suggested but not answered:
-    \n ------- \n
-    {failed_sub_questions}
-    \n ------- \n
-
-    And here are the entities, relationships and terms extracted from the context:
-    \n ------- \n
-    {entity_term_extraction_str}
-    \n ------- \n
-
-   Please generate the list of good, fully contextualized sub-questions that would help to address the
-   main question. Again, please find questions that are NOT overlapping too much with the already answered
-   sub-questions or those that already were suggested and failed.
-   In other words - what can we try in addition to what has been tried so far?
-
-   Please think through it step by step and then generate the list of json dictionaries with the following
-   format:
-
-   {{"sub_questions": [{{"sub_question": <sub-question>,
-        "explanation": <explanation>,
-        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
-        ...]}} """
-
-DEEP_DECOMPOSE_PROMPT = """ \n
-    An initial user question needs to be answered. An initial answer has been provided but it wasn't quite
-    good enough. Also, some sub-questions had been answered and this information has been used to provide
-    the initial answer. Some other subquestions may have been suggested based on little knowledge, but they
-    were not directly answerable. Also, some entities, relationships and terms are givenm to you so that
-    you have an idea of how the avaiolable data looks like.
-
-    Your role is to generate 4-6 new sub-questions that would help to answer the initial question,
-    considering:
-
-    1) The initial question
-    2) The initial answer that was found to be unsatisfactory
-    3) The sub-questions that were answered
-    4) The sub-questions that were suggested but not answered
-    5) The entities, relationships and terms that were extracted from the context
-
-    The individual questions should be answerable by a good RAG system.
-    So a good idea would be to use the sub-questions to resolve ambiguities and/or to separate the
-    question for different entities that may be involved in the original question, but in a way that does
-    not duplicate questions that were already tried.
-
-    Additional Guidelines:
-    - The sub-questions should be specific to the question and provide richer context for the question,
-    resolve ambiguities, or address shortcoming of the initial answer
-    - Each sub-question - when answered - should be relevant for the answer to the original question
-    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
-    other complications that may require extra context.
-    - The sub-questions MUST have the full context of the original question so that it can be executed by
-    a RAG system independently without the original question available
-      (Example:
-        - initial question: "What is the capital of France?"
-        - bad sub-question: "What is the name of the river there?"
-        - good sub-question: "What is the name of the river that flows through Paris?"
-    - For each sub-question, please also provide a search term that can be used to retrieve relevant
-    documents from a document store.
-    \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    Here is the initial sub-optimal answer:
-    \n ------- \n
-    {base_answer}
-    \n ------- \n
-
-    Here are the sub-questions that were answered:
-    \n ------- \n
-    {answered_sub_questions}
-    \n ------- \n
-
-    Here are the sub-questions that were suggested but not answered:
-    \n ------- \n
-    {failed_sub_questions}
-    \n ------- \n
-
-    And here are the entities, relationships and terms extracted from the context:
-    \n ------- \n
-    {entity_term_extraction_str}
-    \n ------- \n
-
-   Please generate the list of good, fully contextualized sub-questions that would help to address the
-   main question. Again, please find questions that are NOT overlapping too much with the already answered
-   sub-questions or those that already were suggested and failed.
-   In other words - what can we try in addition to what has been tried so far?
-
-   Generate the list of json dictionaries with the following format:
-
-   {{"sub_questions": [{{"sub_question": <sub-question>,
-        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
-        ...]}} """
-
-DECOMPOSE_PROMPT = """ \n
-    For an initial user question, please generate at 5-10 individual sub-questions whose answers would help
-    \n to answer the initial question. The individual questions should be answerable by a good RAG system.
-    So a good idea would be to \n use the sub-questions to resolve ambiguities and/or to separate the
-    question for different entities that may be involved in the original question.
-
-    In order to arrive at meaningful sub-questions, please also consider the context retrieved from the
-    document store, expressed as entities, relationships and terms. You can also think about the types
-    mentioned in brackets
-
-    Guidelines:
-    - The sub-questions should be specific to the question and provide richer context for the question,
-    and or resolve ambiguities
-    - Each sub-question - when answered - should be relevant for the answer to the original question
-    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
-    other complications that may require extra context.
-    - The sub-questions MUST have the full context of the original question so that it can be executed by
-    a RAG system independently without the original question available
-      (Example:
-        - initial question: "What is the capital of France?"
-        - bad sub-question: "What is the name of the river there?"
-        - good sub-question: "What is the name of the river that flows through Paris?"
-    - For each sub-question, please provide a short explanation for why it is a good sub-question. So
-    generate a list of dictionaries with the following format:
-      [{{"sub_question": <sub-question>, "explanation": <explanation>, "search_term": <rewrite the
-      sub-question using as a search phrase for the document store>}}, ...]
-
-    \n\n
-    Here is the initial question:
-    \n ------- \n
-    {question}
-    \n ------- \n
-
-    And here are the entities, relationships and terms extracted from the context:
-    \n ------- \n
-    {entity_term_extraction_str}
-    \n ------- \n
-
-   Please generate the list of good, fully contextualized sub-questions that would help to address the
-   main question. Don't be too specific unless the original question is specific.
-   Please think through it step by step and then generate the list of json dictionaries with the following
-   format:
-   {{"sub_questions": [{{"sub_question": <sub-question>,
-        "explanation": <explanation>,
-        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
-        ...]}} """
-
-#### Consolidations
-COMBINED_CONTEXT = """-------
-    Below you will find useful information to answer the original question. First, you see a number of
-    sub-questions with their answers. This information should be considered to be more focussed and
-    somewhat more specific to the original question as it tries to contextualized facts.
-    After that will see the documents that were considered to be relevant to answer the original question.
-
-    Here are the sub-questions and their answers:
-    \n\n {deep_answer_context} \n\n
-    \n\n Here are the documents that were considered to be relevant to answer the original question:
-    \n\n {formated_docs} \n\n
-    ----------------
-    """
-
-SUB_QUESTION_EXPLANATION_RANKER_PROMPT = """-------
-    Below you will find a question that we ultimately want to answer (the original question) and a list of
-    motivations in arbitrary order for generated sub-questions that are supposed to help us answering the
-    original question. The motivations are formatted as <motivation number>:  <motivation explanation>.
-    (Again, the numbering is arbitrary and does not necessarily mean that 1 is the most relevant
-    motivation and 2 is less relevant.)
-
-    Please rank the motivations in order of relevance for answering the original question. Also, try to
-    ensure that the top questions do not duplicate too much, i.e. that they are not too similar.
-    Ultimately, create a list with the motivation numbers where the number of the most relevant
-    motivations comes first.
-
-    Here is the original question:
-    \n\n {original_question} \n\n
-    \n\n Here is the list of sub-question motivations:
-    \n\n {sub_question_explanations} \n\n
-    ----------------
-
-    Please think step by step and then generate the ranked list of motivations.
-
-    Please format your answer as a json object in the following format:
-    {{"reasonning": <explain your reasoning for the ranking>,
-      "ranked_motivations": <ranked list of motivation numbers>}}
-    """

backend/danswer/agent_search/shared_graph_utils/utils.py

@@ -1,91 +0,0 @@
-import ast
-import json
-import re
-from collections.abc import Sequence
-from datetime import datetime
-from datetime import timedelta
-from typing import Any
-
-from danswer.context.search.models import InferenceSection
-
-
-def normalize_whitespace(text: str) -> str:
-    """Normalize whitespace in text to single spaces and strip leading/trailing whitespace."""
-    import re
-
-    return re.sub(r"\s+", " ", text.strip())
-
-
-# Post-processing
-def format_docs(docs: Sequence[InferenceSection]) -> str:
-    return "\n\n".join(doc.combined_content for doc in docs)
-
-
-def clean_and_parse_list_string(json_string: str) -> list[dict]:
-    # Remove markdown code block markers and any newline prefixes
-    cleaned_string = re.sub(r"```json\n|\n```", "", json_string)
-    cleaned_string = cleaned_string.replace("\\n", " ").replace("\n", " ")
-    cleaned_string = " ".join(cleaned_string.split())
-    # Parse the cleaned string into a Python dictionary
-    return ast.literal_eval(cleaned_string)
-
-
-def clean_and_parse_json_string(json_string: str) -> dict[str, Any]:
-    # Remove markdown code block markers and any newline prefixes
-    cleaned_string = re.sub(r"```json\n|\n```", "", json_string)
-    cleaned_string = cleaned_string.replace("\\n", " ").replace("\n", " ")
-    cleaned_string = " ".join(cleaned_string.split())
-    # Parse the cleaned string into a Python dictionary
-    return json.loads(cleaned_string)
-
-
-def format_entity_term_extraction(entity_term_extraction_dict: dict[str, Any]) -> str:
-    entities = entity_term_extraction_dict["entities"]
-    terms = entity_term_extraction_dict["terms"]
-    relationships = entity_term_extraction_dict["relationships"]
-
-    entity_strs = ["\nEntities:\n"]
-    for entity in entities:
-        entity_str = f"{entity['entity_name']} ({entity['entity_type']})"
-        entity_strs.append(entity_str)
-
-    entity_str = "\n - ".join(entity_strs)
-
-    relationship_strs = ["\n\nRelationships:\n"]
-    for relationship in relationships:
-        relationship_str = f"{relationship['name']} ({relationship['type']}): {relationship['entities']}"
-        relationship_strs.append(relationship_str)
-
-    relationship_str = "\n - ".join(relationship_strs)
-
-    term_strs = ["\n\nTerms:\n"]
-    for term in terms:
-        term_str = f"{term['term_name']} ({term['term_type']}): similar to {term['similar_to']}"
-        term_strs.append(term_str)
-
-    term_str = "\n - ".join(term_strs)
-
-    return "\n".join(entity_strs + relationship_strs + term_strs)
-
-
-def _format_time_delta(time: timedelta) -> str:
-    seconds_from_start = f"{((time).seconds):03d}"
-    microseconds_from_start = f"{((time).microseconds):06d}"
-    return f"{seconds_from_start}.{microseconds_from_start}"
-
-
-def generate_log_message(
-    message: str,
-    node_start_time: datetime,
-    graph_start_time: datetime | None = None,
-) -> str:
-    current_time = datetime.now()
-
-    if graph_start_time is not None:
-        graph_time_str = _format_time_delta(current_time - graph_start_time)
-    else:
-        graph_time_str = "N/A"
-
-    node_time_str = _format_time_delta(current_time - node_start_time)
-
-    return f"{graph_time_str} ({node_time_str} s): {message}"

backend/danswer/auth/invited_users.py

@@ -2,8 +2,8 @@ from typing import cast
 
 from danswer.configs.constants import KV_USER_STORE_KEY
 from danswer.key_value_store.factory import get_kv_store
+from danswer.key_value_store.interface import JSON_ro
 from danswer.key_value_store.interface import KvKeyNotFoundError
-from danswer.utils.special_types import JSON_ro
 
 
 def get_invited_users() -> list[str]:

backend/danswer/auth/noauth_user.py

@@ -23,9 +23,7 @@ def load_no_auth_user_preferences(store: KeyValueStore) -> UserPreferences:
         )
         return UserPreferences(**preferences_data)
     except KvKeyNotFoundError:
-        return UserPreferences(
-            chosen_assistants=None, default_model=None, auto_scroll=True
-        )
+        return UserPreferences(chosen_assistants=None, default_model=None)
 
 
 def fetch_no_auth_user(store: KeyValueStore) -> UserInfo:

backend/danswer/auth/schemas.py

@@ -13,24 +13,12 @@ class UserRole(str, Enum):
         groups they are curators of
     - Global Curator can perform admin actions
         for all groups they are a member of
-    - Limited can access a limited set of basic api endpoints
-    - Slack are users that have used danswer via slack but dont have a web login
-    - External permissioned users that have been picked up during the external permissions sync process but don't have a web login
     """
 
-    LIMITED = "limited"
     BASIC = "basic"
     ADMIN = "admin"
     CURATOR = "curator"
     GLOBAL_CURATOR = "global_curator"
-    SLACK_USER = "slack_user"
-    EXT_PERM_USER = "ext_perm_user"
-
-    def is_web_login(self) -> bool:
-        return self not in [
-            UserRole.SLACK_USER,
-            UserRole.EXT_PERM_USER,
-        ]
 
 
 class UserStatus(str, Enum):
@@ -45,8 +33,10 @@ class UserRead(schemas.BaseUser[uuid.UUID]):
 
 class UserCreate(schemas.BaseUserCreate):
     role: UserRole = UserRole.BASIC
+    has_web_login: bool | None = True
     tenant_id: str | None = None
 
 
 class UserUpdate(schemas.BaseUserUpdate):
     role: UserRole
+    has_web_login: bool | None = True

backend/danswer/auth/users.py

@@ -49,7 +49,8 @@ from httpx_oauth.oauth2 import BaseOAuth2
 from httpx_oauth.oauth2 import OAuth2Token
 from pydantic import BaseModel
 from sqlalchemy import text
-from sqlalchemy.ext.asyncio import AsyncSession
+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
@@ -80,14 +81,13 @@ from danswer.db.auth import get_default_admin_user_emails
 from danswer.db.auth import get_user_count
 from danswer.db.auth import get_user_db
 from danswer.db.auth import SQLAlchemyUserAdminDB
-from danswer.db.engine import get_async_session
 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.models import AccessToken
 from danswer.db.models import OAuthAccount
 from danswer.db.models import User
 from danswer.db.users import get_user_by_email
-from danswer.server.utils import BasicAuthenticationError
 from danswer.utils.logger import setup_logger
 from danswer.utils.telemetry import optional_telemetry
 from danswer.utils.telemetry import RecordType
@@ -100,6 +100,11 @@ 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
@@ -217,25 +222,18 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
     reset_password_token_secret = USER_AUTH_SECRET
     verification_token_secret = USER_AUTH_SECRET
 
-    user_db: SQLAlchemyUserDatabase[User, uuid.UUID]
-
     async def create(
         self,
         user_create: schemas.UC | UserCreate,
         safe: bool = False,
         request: Optional[Request] = None,
     ) -> User:
-        referral_source = None
-        if request is not None:
-            referral_source = request.cookies.get("referral_source", None)
-
         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,
-            referral_source=referral_source,
         )
 
         async with get_async_session_with_tenant(tenant_id) as db_session:
@@ -244,9 +242,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             verify_email_is_invited(user_create.email)
             verify_email_domain(user_create.email)
             if MULTI_TENANT:
-                tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](
-                    db_session, User, OAuthAccount
-                )
+                tenant_user_db = SQLAlchemyUserAdminDB(db_session, User, OAuthAccount)
                 self.user_db = tenant_user_db
                 self.database = tenant_user_db
 
@@ -265,9 +261,14 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             except exceptions.UserAlreadyExists:
                 user = await self.get_by_email(user_create.email)
                 # Handle case where user has used product outside of web and is now creating an account through web
-                if not user.role.is_web_login() and user_create.role.is_web_login():
+                if (
+                    not user.has_web_login
+                    and hasattr(user_create, "has_web_login")
+                    and user_create.has_web_login
+                ):
                     user_update = UserUpdate(
                         password=user_create.password,
+                        has_web_login=True,
                         role=user_create.role,
                         is_verified=user_create.is_verified,
                     )
@@ -281,7 +282,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             return user
 
     async def oauth_callback(
-        self,
+        self: "BaseUserManager[models.UOAP, models.ID]",
         oauth_name: str,
         access_token: str,
         account_id: str,
@@ -292,18 +293,13 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
         *,
         associate_by_email: bool = False,
         is_verified_by_default: bool = False,
-    ) -> User:
-        referral_source = None
-        if request:
-            referral_source = getattr(request.state, "referral_source", None)
-
+    ) -> models.UOAP:
         tenant_id = await fetch_ee_implementation_or_noop(
             "danswer.server.tenants.provisioning",
             "get_or_create_tenant_id",
             async_return_default_schema,
         )(
             email=account_email,
-            referral_source=referral_source,
         )
 
         if not tenant_id:
@@ -318,11 +314,9 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             verify_email_domain(account_email)
 
             if MULTI_TENANT:
-                tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](
-                    db_session, User, OAuthAccount
-                )
+                tenant_user_db = SQLAlchemyUserAdminDB(db_session, User, OAuthAccount)
                 self.user_db = tenant_user_db
-                self.database = tenant_user_db
+                self.database = tenant_user_db  # type: ignore
 
             oauth_account_dict = {
                 "oauth_name": oauth_name,
@@ -374,11 +368,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                         and existing_oauth_account.oauth_name == oauth_name
                     ):
                         user = await self.user_db.update_oauth_account(
-                            user,
-                            # NOTE: OAuthAccount DOES implement the OAuthAccountProtocol
-                            # but the type checker doesn't know that :(
-                            existing_oauth_account,  # type: ignore
-                            oauth_account_dict,
+                            user, existing_oauth_account, oauth_account_dict
                         )
 
             # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to
@@ -391,15 +381,16 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                 )
 
             # Handle case where user has used product outside of web and is now creating an account through web
-            if not user.role.is_web_login():
+            if not user.has_web_login:  # type: ignore
                 await self.user_db.update(
                     user,
                     {
                         "is_verified": is_verified_by_default,
-                        "role": UserRole.BASIC,
+                        "has_web_login": True,
                     },
                 )
                 user.is_verified = is_verified_by_default
+                user.has_web_login = True  # type: ignore
 
             # this is needed if an organization goes from `TRACK_EXTERNAL_IDP_EXPIRY=true` to `false`
             # otherwise, the oidc expiry will always be old, and the user will never be able to login
@@ -474,7 +465,9 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                 self.password_helper.hash(credentials.password)
                 return None
 
-            if not user.role.is_web_login():
+            has_web_login = attributes.get_attribute(user, "has_web_login")
+
+            if not has_web_login:
                 raise BasicAuthenticationError(
                     detail="NO_WEB_LOGIN_AND_HAS_NO_PASSWORD",
                 )
@@ -605,7 +598,7 @@ optional_fastapi_current_user = fastapi_users.current_user(active=True, optional
 async def optional_user_(
     request: Request,
     user: User | None,
-    async_db_session: AsyncSession,
+    db_session: Session,
 ) -> User | None:
     """NOTE: `request` and `db_session` are not used here, but are included
     for the EE version of this function."""
@@ -614,21 +607,13 @@ async def optional_user_(
 
 async def optional_user(
     request: Request,
-    async_db_session: AsyncSession = Depends(get_async_session),
+    db_session: Session = Depends(get_session),
     user: User | None = Depends(optional_fastapi_current_user),
 ) -> User | None:
     versioned_fetch_user = fetch_versioned_implementation(
         "danswer.auth.users", "optional_user_"
     )
-    user = await versioned_fetch_user(request, user, async_db_session)
-
-    # check if an API key is present
-    if user is None:
-        hashed_api_key = get_hashed_api_key_from_request(request)
-        if hashed_api_key:
-            user = await fetch_user_for_api_key(hashed_api_key, async_db_session)
-
-    return user
+    return await versioned_fetch_user(request, user, db_session)
 
 
 async def double_check_user(
@@ -667,24 +652,10 @@ async def current_user_with_expired_token(
     return await double_check_user(user, include_expired=True)
 
 
-async def current_limited_user(
-    user: User | None = Depends(optional_user),
-) -> User | None:
-    return await double_check_user(user)
-
-
 async def current_user(
     user: User | None = Depends(optional_user),
 ) -> User | None:
-    user = await double_check_user(user)
-    if not user:
-        return None
-
-    if user.role == UserRole.LIMITED:
-        raise BasicAuthenticationError(
-            detail="Access denied. User role is LIMITED. BASIC or higher permissions are required.",
-        )
-    return user
+    return await double_check_user(user)
 
 
 async def current_curator_or_admin_user(
@@ -740,6 +711,8 @@ def generate_state_token(
 
 
 # refer to https://github.com/fastapi-users/fastapi-users/blob/42ddc241b965475390e2bce887b084152ae1a2cd/fastapi_users/fastapi_users.py#L91
+
+
 def create_danswer_oauth_router(
     oauth_client: BaseOAuth2,
     backend: AuthenticationBackend,
@@ -789,22 +762,15 @@ def get_oauth_router(
         response_model=OAuth2AuthorizeResponse,
     )
     async def authorize(
-        request: Request,
-        scopes: List[str] = Query(None),
+        request: Request, scopes: List[str] = Query(None)
     ) -> OAuth2AuthorizeResponse:
-        referral_source = request.cookies.get("referral_source", None)
-
         if redirect_url is not None:
             authorize_redirect_url = redirect_url
         else:
             authorize_redirect_url = str(request.url_for(callback_route_name))
 
         next_url = request.query_params.get("next", "/")
-
-        state_data: Dict[str, str] = {
-            "next_url": next_url,
-            "referral_source": referral_source or "default_referral",
-        }
+        state_data: Dict[str, str] = {"next_url": next_url}
         state = generate_state_token(state_data, state_secret)
         authorization_url = await oauth_client.get_authorization_url(
             authorize_redirect_url,
@@ -863,11 +829,8 @@ def get_oauth_router(
             raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST)
 
         next_url = state_data.get("next_url", "/")
-        referral_source = state_data.get("referral_source", None)
 
-        request.state.referral_source = referral_source
-
-        # Proceed to authenticate or create the user
+        # Authenticate user
         try:
             user = await user_manager.oauth_callback(
                 oauth_client.name,
@@ -909,13 +872,14 @@ def get_oauth_router(
             redirect_response.status_code = response.status_code
         if hasattr(response, "media_type"):
             redirect_response.media_type = response.media_type
+
         return redirect_response
 
     return router
 
 
-async def api_key_dep(
-    request: Request, async_db_session: AsyncSession = Depends(get_async_session)
+def api_key_dep(
+    request: Request, db_session: Session = Depends(get_session)
 ) -> User | None:
     if AUTH_TYPE == AuthType.DISABLED:
         return None
@@ -925,7 +889,7 @@ async def api_key_dep(
         raise HTTPException(status_code=401, detail="Missing API key")
 
     if hashed_api_key:
-        user = await fetch_user_for_api_key(hashed_api_key, async_db_session)
+        user = fetch_user_for_api_key(hashed_api_key, db_session)
 
     if user is None:
         raise HTTPException(status_code=401, detail="Invalid API key")

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

@@ -11,7 +11,6 @@ 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 redis.lock import Lock as RedisLock
 from sentry_sdk.integrations.celery import CeleryIntegration
 from sqlalchemy import text
 from sqlalchemy.orm import Session
@@ -25,8 +24,6 @@ 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_doc_perm_sync import RedisConnectorPermissionSync
-from danswer.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync
 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
@@ -139,22 +136,6 @@ def on_task_postrun(
             RedisConnectorPrune.remove_from_taskset(int(cc_pair_id), task_id, r)
         return
 
-    if task_id.startswith(RedisConnectorPermissionSync.SUBTASK_PREFIX):
-        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
-        if cc_pair_id is not None:
-            RedisConnectorPermissionSync.remove_from_taskset(
-                int(cc_pair_id), task_id, r
-            )
-        return
-
-    if task_id.startswith(RedisConnectorExternalGroupSync.SUBTASK_PREFIX):
-        cc_pair_id = RedisConnector.get_id_from_task_id(task_id)
-        if cc_pair_id is not None:
-            RedisConnectorExternalGroupSync.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"""
@@ -333,16 +314,16 @@ def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
         return
 
     logger.info("Releasing primary worker lock.")
-    lock: RedisLock = sender.primary_worker_lock
+    lock = sender.primary_worker_lock
     try:
         if lock.owned():
             try:
                 lock.release()
                 sender.primary_worker_lock = None
-            except Exception:
-                logger.exception("Failed to release primary worker lock")
-    except Exception:
-        logger.exception("Failed to check if primary worker lock is owned")
+            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(

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

@@ -12,7 +12,6 @@ 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 IGNORED_SYNCING_TENANT_LIST
 from shared_configs.configs import MULTI_TENANT
 
 logger = setup_logger(__name__)
@@ -73,15 +72,6 @@ class DynamicTenantScheduler(PersistentScheduler):
             logger.info(f"Found {len(existing_tenants)} existing tenants in schedule")
 
             for tenant_id in tenant_ids:
-                if (
-                    IGNORED_SYNCING_TENANT_LIST
-                    and tenant_id in IGNORED_SYNCING_TENANT_LIST
-                ):
-                    logger.info(
-                        f"Skipping tenant {tenant_id} as it is in the ignored syncing list"
-                    )
-                    continue
-
                 if tenant_id not in existing_tenants:
                     logger.info(f"Processing new tenant: {tenant_id}")
 

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

@@ -91,7 +91,5 @@ def on_setup_logging(
 celery_app.autodiscover_tasks(
     [
         "danswer.background.celery.tasks.pruning",
-        "danswer.background.celery.tasks.doc_permission_syncing",
-        "danswer.background.celery.tasks.external_group_syncing",
     ]
 )

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

@@ -6,7 +6,6 @@ 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_process_init
 from celery.signals import worker_ready
 from celery.signals import worker_shutdown
 
@@ -60,7 +59,7 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
     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=sender.concurrency, max_overflow=sender.concurrency)
+    SqlEngine.init_engine(pool_size=8, max_overflow=0)
 
     # Startup checks are not needed in multi-tenant case
     if MULTI_TENANT:
@@ -82,11 +81,6 @@ def on_worker_shutdown(sender: Any, **kwargs: Any) -> None:
     app_base.on_worker_shutdown(sender, **kwargs)
 
 
-@worker_process_init.connect
-def init_worker(**kwargs: Any) -> None:
-    SqlEngine.reset_engine()
-
-
 @signals.setup_logging.connect
 def on_setup_logging(
     loglevel: Any, logfile: Any, format: Any, colorize: Any, **kwargs: Any

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

@@ -92,6 +92,5 @@ celery_app.autodiscover_tasks(
         "danswer.background.celery.tasks.shared",
         "danswer.background.celery.tasks.vespa",
         "danswer.background.celery.tasks.connector_deletion",
-        "danswer.background.celery.tasks.doc_permission_syncing",
     ]
 )

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

@@ -1,6 +1,5 @@
 import multiprocessing
 from typing import Any
-from typing import cast
 
 from celery import bootsteps  # type: ignore
 from celery import Celery
@@ -11,25 +10,16 @@ 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 redis.lock import Lock as RedisLock
 
 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.background.celery.tasks.indexing.tasks import (
-    get_unfenced_index_attempt_ids,
-)
 from danswer.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME
-from danswer.db.engine import get_session_with_default_tenant
 from danswer.db.engine import SqlEngine
-from danswer.db.index_attempt import get_index_attempt
-from danswer.db.index_attempt import mark_attempt_canceled
 from danswer.redis.redis_connector_credential_pair import RedisConnectorCredentialPair
 from danswer.redis.redis_connector_delete import RedisConnectorDelete
-from danswer.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync
-from danswer.redis.redis_connector_ext_group_sync import RedisConnectorExternalGroupSync
 from danswer.redis.redis_connector_index import RedisConnectorIndex
 from danswer.redis.redis_connector_prune import RedisConnectorPrune
 from danswer.redis.redis_connector_stop import RedisConnectorStop
@@ -39,6 +29,7 @@ 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__)
@@ -98,15 +89,6 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
     # by the primary worker. This is unnecessary in the multi tenant scenario
     r = get_redis_client(tenant_id=None)
 
-    # Log the role and slave count - being connected to a slave or slave count > 0 could be problematic
-    info: dict[str, Any] = cast(dict, r.info("replication"))
-    role: str = cast(str, info.get("role"))
-    connected_slaves: int = info.get("connected_slaves", 0)
-
-    logger.info(
-        f"Redis INFO REPLICATION: role={role} connected_slaves={connected_slaves}"
-    )
-
     # 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
@@ -116,13 +98,9 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
     # 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.
-
-    # set thread_local=False since we don't control what thread the periodic task might
-    # reacquire the lock with
-    lock: RedisLock = r.lock(
+    lock = r.lock(
         DanswerRedisLocks.PRIMARY_WORKER,
         timeout=CELERY_PRIMARY_WORKER_LOCK_TIMEOUT,
-        thread_local=False,
     )
 
     logger.info("Primary worker lock: Acquire starting.")
@@ -156,27 +134,6 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
 
     RedisConnectorStop.reset_all(r)
 
-    RedisConnectorPermissionSync.reset_all(r)
-
-    RedisConnectorExternalGroupSync.reset_all(r)
-
-    # mark orphaned index attempts as failed
-    with get_session_with_default_tenant() as db_session:
-        unfenced_attempt_ids = get_unfenced_index_attempt_ids(db_session, r)
-        for attempt_id in unfenced_attempt_ids:
-            attempt = get_index_attempt(db_session, attempt_id)
-            if not attempt:
-                continue
-
-            failure_reason = (
-                f"Canceling leftover index attempt found on startup: "
-                f"index_attempt={attempt.id} "
-                f"cc_pair={attempt.connector_credential_pair_id} "
-                f"search_settings={attempt.search_settings_id}"
-            )
-            logger.warning(failure_reason)
-            mark_attempt_canceled(attempt.id, db_session, failure_reason)
-
 
 @worker_ready.connect
 def on_worker_ready(sender: Any, **kwargs: Any) -> None:
@@ -231,7 +188,7 @@ class HubPeriodicTask(bootsteps.StartStopStep):
             if not hasattr(worker, "primary_worker_lock"):
                 return
 
-            lock: RedisLock = worker.primary_worker_lock
+            lock = worker.primary_worker_lock
 
             r = get_redis_client(tenant_id=None)
 
@@ -276,8 +233,6 @@ celery_app.autodiscover_tasks(
         "danswer.background.celery.tasks.connector_deletion",
         "danswer.background.celery.tasks.indexing",
         "danswer.background.celery.tasks.periodic",
-        "danswer.background.celery.tasks.doc_permission_syncing",
-        "danswer.background.celery.tasks.external_group_syncing",
         "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/celery_utils.py

@@ -4,6 +4,7 @@ from typing import Any
 
 from sqlalchemy.orm import Session
 
+from danswer.background.indexing.run_indexing import RunIndexingCallbackInterface
 from danswer.configs.app_configs import MAX_PRUNING_DOCUMENT_RETRIEVAL_PER_MINUTE
 from danswer.connectors.cross_connector_utils.rate_limit_wrapper import (
     rate_limit_builder,
@@ -16,7 +17,6 @@ from danswer.connectors.models import Document
 from danswer.db.connector_credential_pair import get_connector_credential_pair
 from danswer.db.enums import TaskStatus
 from danswer.db.models import TaskQueueState
-from danswer.indexing.indexing_heartbeat import IndexingHeartbeatInterface
 from danswer.redis.redis_connector import RedisConnector
 from danswer.server.documents.models import DeletionAttemptSnapshot
 from danswer.utils.logger import setup_logger
@@ -78,10 +78,10 @@ def document_batch_to_ids(
 
 def extract_ids_from_runnable_connector(
     runnable_connector: BaseConnector,
-    callback: IndexingHeartbeatInterface | None = None,
+    callback: RunIndexingCallbackInterface | None = None,
 ) -> set[str]:
     """
-    If the SlimConnector hasnt been implemented for the given connector, just pull
+    If the PruneConnector hasnt been implemented for the given connector, just pull
     all docs using the load_from_state and grab out the IDs.
 
     Optionally, a callback can be passed to handle the length of each document batch.
@@ -111,15 +111,10 @@ def extract_ids_from_runnable_connector(
     for doc_batch in doc_batch_generator:
         if callback:
             if callback.should_stop():
-                raise RuntimeError(
-                    "extract_ids_from_runnable_connector: Stop signal detected"
-                )
-
+                raise RuntimeError("Stop signal received")
+            callback.progress(len(doc_batch))
         all_connector_doc_ids.update(doc_batch_processing_func(doc_batch))
 
-        if callback:
-            callback.progress("extract_ids_from_runnable_connector", len(doc_batch))
-
     return all_connector_doc_ids
 
 

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

@@ -2,58 +2,45 @@ from datetime import timedelta
 from typing import Any
 
 from danswer.configs.constants import DanswerCeleryPriority
-from danswer.configs.constants import DanswerCeleryTask
 
 
 tasks_to_schedule = [
     {
         "name": "check-for-vespa-sync",
-        "task": DanswerCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
-        "schedule": timedelta(seconds=20),
+        "task": "check_for_vespa_sync_task",
+        "schedule": timedelta(seconds=5),
         "options": {"priority": DanswerCeleryPriority.HIGH},
     },
     {
         "name": "check-for-connector-deletion",
-        "task": DanswerCeleryTask.CHECK_FOR_CONNECTOR_DELETION,
+        "task": "check_for_connector_deletion_task",
         "schedule": timedelta(seconds=20),
         "options": {"priority": DanswerCeleryPriority.HIGH},
     },
     {
         "name": "check-for-indexing",
-        "task": DanswerCeleryTask.CHECK_FOR_INDEXING,
-        "schedule": timedelta(seconds=15),
+        "task": "check_for_indexing",
+        "schedule": timedelta(seconds=10),
         "options": {"priority": DanswerCeleryPriority.HIGH},
     },
     {
         "name": "check-for-prune",
-        "task": DanswerCeleryTask.CHECK_FOR_PRUNING,
-        "schedule": timedelta(seconds=15),
+        "task": "check_for_pruning",
+        "schedule": timedelta(seconds=10),
         "options": {"priority": DanswerCeleryPriority.HIGH},
     },
     {
         "name": "kombu-message-cleanup",
-        "task": DanswerCeleryTask.KOMBU_MESSAGE_CLEANUP_TASK,
+        "task": "kombu_message_cleanup_task",
         "schedule": timedelta(seconds=3600),
         "options": {"priority": DanswerCeleryPriority.LOWEST},
     },
     {
         "name": "monitor-vespa-sync",
-        "task": DanswerCeleryTask.MONITOR_VESPA_SYNC,
+        "task": "monitor_vespa_sync",
         "schedule": timedelta(seconds=5),
         "options": {"priority": DanswerCeleryPriority.HIGH},
     },
-    {
-        "name": "check-for-doc-permissions-sync",
-        "task": DanswerCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
-        "schedule": timedelta(seconds=30),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
-    {
-        "name": "check-for-external-group-sync",
-        "task": DanswerCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,
-        "schedule": timedelta(seconds=20),
-        "options": {"priority": DanswerCeleryPriority.HIGH},
-    },
 ]
 
 

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

@@ -1,17 +1,17 @@
 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.lock import Lock as RedisLock
+from redis import Redis
 from sqlalchemy.orm import Session
 
 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 DanswerCeleryTask
 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
@@ -19,7 +19,7 @@ from danswer.db.engine import get_session_with_tenant
 from danswer.db.enums import ConnectorCredentialPairStatus
 from danswer.db.search_settings import get_all_search_settings
 from danswer.redis.redis_connector import RedisConnector
-from danswer.redis.redis_connector_delete import RedisConnectorDeletePayload
+from danswer.redis.redis_connector_delete import RedisConnectorDeletionFenceData
 from danswer.redis.redis_pool import get_redis_client
 
 
@@ -29,7 +29,7 @@ class TaskDependencyError(RuntimeError):
 
 
 @shared_task(
-    name=DanswerCeleryTask.CHECK_FOR_CONNECTOR_DELETION,
+    name="check_for_connector_deletion_task",
     soft_time_limit=JOB_TIMEOUT,
     trail=False,
     bind=True,
@@ -37,7 +37,7 @@ class TaskDependencyError(RuntimeError):
 def check_for_connector_deletion_task(self: Task, *, tenant_id: str | None) -> None:
     r = get_redis_client(tenant_id=tenant_id)
 
-    lock_beat: RedisLock = r.lock(
+    lock_beat = r.lock(
         DanswerRedisLocks.CHECK_CONNECTOR_DELETION_BEAT_LOCK,
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
@@ -60,7 +60,7 @@ def check_for_connector_deletion_task(self: Task, *, tenant_id: str | None) -> N
                 redis_connector = RedisConnector(tenant_id, cc_pair_id)
                 try:
                     try_generate_document_cc_pair_cleanup_tasks(
-                        self.app, cc_pair_id, db_session, lock_beat, tenant_id
+                        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
@@ -86,7 +86,8 @@ def try_generate_document_cc_pair_cleanup_tasks(
     app: Celery,
     cc_pair_id: int,
     db_session: Session,
-    lock_beat: RedisLock,
+    r: Redis,
+    lock_beat: redis.lock.Lock,
     tenant_id: str | None,
 ) -> int | None:
     """Returns an int if syncing is needed. The int represents the number of sync tasks generated.
@@ -117,7 +118,7 @@ def try_generate_document_cc_pair_cleanup_tasks(
         return None
 
     # set a basic fence to start
-    fence_payload = RedisConnectorDeletePayload(
+    fence_payload = RedisConnectorDeletionFenceData(
         num_tasks=None,
         submitted=datetime.now(timezone.utc),
     )
@@ -142,12 +143,6 @@ def try_generate_document_cc_pair_cleanup_tasks(
                 f"cc_pair={cc_pair_id}"
             )
 
-        if redis_connector.permissions.fenced:
-            raise TaskDependencyError(
-                f"Connector deletion - Delayed (permissions 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()
 

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

@@ -1,345 +0,0 @@
-from datetime import datetime
-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 redis.lock import Lock as RedisLock
-
-from danswer.access.models import DocExternalAccess
-from danswer.background.celery.apps.app_base import task_logger
-from danswer.configs.app_configs import JOB_TIMEOUT
-from danswer.configs.constants import CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT
-from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
-from danswer.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX
-from danswer.configs.constants import DanswerCeleryPriority
-from danswer.configs.constants import DanswerCeleryQueues
-from danswer.configs.constants import DanswerCeleryTask
-from danswer.configs.constants import DanswerRedisLocks
-from danswer.configs.constants import DocumentSource
-from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
-from danswer.db.document import upsert_document_by_connector_credential_pair
-from danswer.db.engine import get_session_with_tenant
-from danswer.db.enums import AccessType
-from danswer.db.enums import ConnectorCredentialPairStatus
-from danswer.db.models import ConnectorCredentialPair
-from danswer.db.users import batch_add_ext_perm_user_if_not_exists
-from danswer.redis.redis_connector import RedisConnector
-from danswer.redis.redis_connector_doc_perm_sync import (
-    RedisConnectorPermissionSyncPayload,
-)
-from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import doc_permission_sync_ctx
-from danswer.utils.logger import setup_logger
-from ee.danswer.db.connector_credential_pair import get_all_auto_sync_cc_pairs
-from ee.danswer.db.document import upsert_document_external_perms
-from ee.danswer.external_permissions.sync_params import DOC_PERMISSION_SYNC_PERIODS
-from ee.danswer.external_permissions.sync_params import DOC_PERMISSIONS_FUNC_MAP
-
-logger = setup_logger()
-
-
-DOCUMENT_PERMISSIONS_UPDATE_MAX_RETRIES = 3
-
-
-# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT
-LIGHT_SOFT_TIME_LIMIT = 105
-LIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15
-
-
-def _is_external_doc_permissions_sync_due(cc_pair: ConnectorCredentialPair) -> bool:
-    """Returns boolean indicating if external doc permissions sync is due."""
-
-    if cc_pair.access_type != AccessType.SYNC:
-        return False
-
-    # skip doc permissions sync if not active
-    if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:
-        return False
-
-    if cc_pair.status == ConnectorCredentialPairStatus.DELETING:
-        return False
-
-    # If the last sync is None, it has never been run so we run the sync
-    last_perm_sync = cc_pair.last_time_perm_sync
-    if last_perm_sync is None:
-        return True
-
-    source_sync_period = DOC_PERMISSION_SYNC_PERIODS.get(cc_pair.connector.source)
-
-    # If RESTRICTED_FETCH_PERIOD[source] is None, we always run the sync.
-    if not source_sync_period:
-        return True
-
-    # If the last sync is greater than the full fetch period, we run the sync
-    next_sync = last_perm_sync + timedelta(seconds=source_sync_period)
-    if datetime.now(timezone.utc) >= next_sync:
-        return True
-
-    return False
-
-
-@shared_task(
-    name=DanswerCeleryTask.CHECK_FOR_DOC_PERMISSIONS_SYNC,
-    soft_time_limit=JOB_TIMEOUT,
-    bind=True,
-)
-def check_for_doc_permissions_sync(self: Task, *, tenant_id: str | None) -> None:
-    r = get_redis_client(tenant_id=tenant_id)
-
-    lock_beat = r.lock(
-        DanswerRedisLocks.CHECK_CONNECTOR_DOC_PERMISSIONS_SYNC_BEAT_LOCK,
-        timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
-    )
-
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return
-
-        # get all cc pairs that need to be synced
-        cc_pair_ids_to_sync: list[int] = []
-        with get_session_with_tenant(tenant_id) as db_session:
-            cc_pairs = get_all_auto_sync_cc_pairs(db_session)
-
-            for cc_pair in cc_pairs:
-                if _is_external_doc_permissions_sync_due(cc_pair):
-                    cc_pair_ids_to_sync.append(cc_pair.id)
-
-        for cc_pair_id in cc_pair_ids_to_sync:
-            tasks_created = try_creating_permissions_sync_task(
-                self.app, cc_pair_id, r, tenant_id
-            )
-            if not tasks_created:
-                continue
-
-            task_logger.info(f"Doc permissions sync 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(f"Unexpected exception: tenant={tenant_id}")
-    finally:
-        if lock_beat.owned():
-            lock_beat.release()
-
-
-def try_creating_permissions_sync_task(
-    app: Celery,
-    cc_pair_id: int,
-    r: Redis,
-    tenant_id: str | None,
-) -> int | None:
-    """Returns an int if syncing is needed. The int represents the number of sync tasks generated.
-    Returns None if no syncing is required."""
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-
-    LOCK_TIMEOUT = 30
-
-    lock: RedisLock = r.lock(
-        DANSWER_REDIS_FUNCTION_LOCK_PREFIX + "try_generate_permissions_sync_tasks",
-        timeout=LOCK_TIMEOUT,
-    )
-
-    acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)
-    if not acquired:
-        return None
-
-    try:
-        if redis_connector.permissions.fenced:
-            return None
-
-        if redis_connector.delete.fenced:
-            return None
-
-        if redis_connector.prune.fenced:
-            return None
-
-        redis_connector.permissions.generator_clear()
-        redis_connector.permissions.taskset_clear()
-
-        custom_task_id = f"{redis_connector.permissions.generator_task_key}_{uuid4()}"
-
-        result = app.send_task(
-            DanswerCeleryTask.CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK,
-            kwargs=dict(
-                cc_pair_id=cc_pair_id,
-                tenant_id=tenant_id,
-            ),
-            queue=DanswerCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC,
-            task_id=custom_task_id,
-            priority=DanswerCeleryPriority.HIGH,
-        )
-
-        # set a basic fence to start
-        payload = RedisConnectorPermissionSyncPayload(
-            started=None, celery_task_id=result.id
-        )
-
-        redis_connector.permissions.set_fence(payload)
-    except Exception:
-        task_logger.exception(f"Unexpected exception: cc_pair={cc_pair_id}")
-        return None
-    finally:
-        if lock.owned():
-            lock.release()
-
-    return 1
-
-
-@shared_task(
-    name=DanswerCeleryTask.CONNECTOR_PERMISSION_SYNC_GENERATOR_TASK,
-    acks_late=False,
-    soft_time_limit=JOB_TIMEOUT,
-    track_started=True,
-    trail=False,
-    bind=True,
-)
-def connector_permission_sync_generator_task(
-    self: Task,
-    cc_pair_id: int,
-    tenant_id: str | None,
-) -> None:
-    """
-    Permission sync task that handles document permission syncing for a given connector credential pair
-    This task assumes that the task has already been properly fenced
-    """
-
-    doc_permission_sync_ctx_dict = doc_permission_sync_ctx.get()
-    doc_permission_sync_ctx_dict["cc_pair_id"] = cc_pair_id
-    doc_permission_sync_ctx_dict["request_id"] = self.request.id
-    doc_permission_sync_ctx.set(doc_permission_sync_ctx_dict)
-
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-
-    r = get_redis_client(tenant_id=tenant_id)
-
-    lock = r.lock(
-        DanswerRedisLocks.CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX
-        + f"_{redis_connector.id}",
-        timeout=CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT,
-    )
-
-    acquired = lock.acquire(blocking=False)
-    if not acquired:
-        task_logger.warning(
-            f"Permission sync task already running, exiting...: cc_pair={cc_pair_id}"
-        )
-        return None
-
-    try:
-        with get_session_with_tenant(tenant_id) as db_session:
-            cc_pair = get_connector_credential_pair_from_id(cc_pair_id, db_session)
-            if cc_pair is None:
-                raise ValueError(
-                    f"No connector credential pair found for id: {cc_pair_id}"
-                )
-
-            source_type = cc_pair.connector.source
-
-            doc_sync_func = DOC_PERMISSIONS_FUNC_MAP.get(source_type)
-            if doc_sync_func is None:
-                raise ValueError(
-                    f"No doc sync func found for {source_type} with cc_pair={cc_pair_id}"
-                )
-
-            logger.info(f"Syncing docs for {source_type} with cc_pair={cc_pair_id}")
-
-            payload = redis_connector.permissions.payload
-            if not payload:
-                raise ValueError(f"No fence payload found: cc_pair={cc_pair_id}")
-
-            payload.started = datetime.now(timezone.utc)
-            redis_connector.permissions.set_fence(payload)
-
-            document_external_accesses: list[DocExternalAccess] = doc_sync_func(cc_pair)
-
-            task_logger.info(
-                f"RedisConnector.permissions.generate_tasks starting. cc_pair={cc_pair_id}"
-            )
-            tasks_generated = redis_connector.permissions.generate_tasks(
-                celery_app=self.app,
-                lock=lock,
-                new_permissions=document_external_accesses,
-                source_string=source_type,
-                connector_id=cc_pair.connector.id,
-                credential_id=cc_pair.credential.id,
-            )
-            if tasks_generated is None:
-                return None
-
-            task_logger.info(
-                f"RedisConnector.permissions.generate_tasks finished. "
-                f"cc_pair={cc_pair_id} tasks_generated={tasks_generated}"
-            )
-
-            redis_connector.permissions.generator_complete = tasks_generated
-
-    except Exception as e:
-        task_logger.exception(f"Failed to run permission sync: cc_pair={cc_pair_id}")
-
-        redis_connector.permissions.generator_clear()
-        redis_connector.permissions.taskset_clear()
-        redis_connector.permissions.set_fence(None)
-        raise e
-    finally:
-        if lock.owned():
-            lock.release()
-
-
-@shared_task(
-    name=DanswerCeleryTask.UPDATE_EXTERNAL_DOCUMENT_PERMISSIONS_TASK,
-    soft_time_limit=LIGHT_SOFT_TIME_LIMIT,
-    time_limit=LIGHT_TIME_LIMIT,
-    max_retries=DOCUMENT_PERMISSIONS_UPDATE_MAX_RETRIES,
-    bind=True,
-)
-def update_external_document_permissions_task(
-    self: Task,
-    tenant_id: str | None,
-    serialized_doc_external_access: dict,
-    source_string: str,
-    connector_id: int,
-    credential_id: int,
-) -> bool:
-    document_external_access = DocExternalAccess.from_dict(
-        serialized_doc_external_access
-    )
-    doc_id = document_external_access.doc_id
-    external_access = document_external_access.external_access
-    try:
-        with get_session_with_tenant(tenant_id) as db_session:
-            # Add the users to the DB if they don't exist
-            batch_add_ext_perm_user_if_not_exists(
-                db_session=db_session,
-                emails=list(external_access.external_user_emails),
-            )
-            # Then we upsert the document's external permissions in postgres
-            created_new_doc = upsert_document_external_perms(
-                db_session=db_session,
-                doc_id=doc_id,
-                external_access=external_access,
-                source_type=DocumentSource(source_string),
-            )
-
-            if created_new_doc:
-                # If a new document was created, we associate it with the cc_pair
-                upsert_document_by_connector_credential_pair(
-                    db_session=db_session,
-                    connector_id=connector_id,
-                    credential_id=credential_id,
-                    document_ids=[doc_id],
-                )
-
-            logger.debug(
-                f"Successfully synced postgres document permissions for {doc_id}"
-            )
-        return True
-    except Exception:
-        logger.exception("Error Syncing Document Permissions")
-        return False

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

@@ -1,298 +0,0 @@
-from datetime import datetime
-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 redis.lock import Lock as RedisLock
-
-from danswer.background.celery.apps.app_base import task_logger
-from danswer.configs.app_configs import JOB_TIMEOUT
-from danswer.configs.constants import CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT
-from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
-from danswer.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX
-from danswer.configs.constants import DanswerCeleryPriority
-from danswer.configs.constants import DanswerCeleryQueues
-from danswer.configs.constants import DanswerCeleryTask
-from danswer.configs.constants import DanswerRedisLocks
-from danswer.db.connector import mark_cc_pair_as_external_group_synced
-from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
-from danswer.db.engine import get_session_with_tenant
-from danswer.db.enums import AccessType
-from danswer.db.enums import ConnectorCredentialPairStatus
-from danswer.db.models import ConnectorCredentialPair
-from danswer.redis.redis_connector import RedisConnector
-from danswer.redis.redis_connector_ext_group_sync import (
-    RedisConnectorExternalGroupSyncPayload,
-)
-from danswer.redis.redis_pool import get_redis_client
-from danswer.utils.logger import setup_logger
-from ee.danswer.db.connector_credential_pair import get_all_auto_sync_cc_pairs
-from ee.danswer.db.connector_credential_pair import get_cc_pairs_by_source
-from ee.danswer.db.external_perm import ExternalUserGroup
-from ee.danswer.db.external_perm import replace_user__ext_group_for_cc_pair
-from ee.danswer.external_permissions.sync_params import EXTERNAL_GROUP_SYNC_PERIODS
-from ee.danswer.external_permissions.sync_params import GROUP_PERMISSIONS_FUNC_MAP
-from ee.danswer.external_permissions.sync_params import (
-    GROUP_PERMISSIONS_IS_CC_PAIR_AGNOSTIC,
-)
-
-logger = setup_logger()
-
-
-EXTERNAL_GROUPS_UPDATE_MAX_RETRIES = 3
-
-
-# 5 seconds more than RetryDocumentIndex STOP_AFTER+MAX_WAIT
-LIGHT_SOFT_TIME_LIMIT = 105
-LIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15
-
-
-def _is_external_group_sync_due(cc_pair: ConnectorCredentialPair) -> bool:
-    """Returns boolean indicating if external group sync is due."""
-
-    if cc_pair.access_type != AccessType.SYNC:
-        return False
-
-    # skip external group sync if not active
-    if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:
-        return False
-
-    if cc_pair.status == ConnectorCredentialPairStatus.DELETING:
-        return False
-
-    # If there is not group sync function for the connector, we don't run the sync
-    # This is fine because all sources dont necessarily have a concept of groups
-    if not GROUP_PERMISSIONS_FUNC_MAP.get(cc_pair.connector.source):
-        return False
-
-    # If the last sync is None, it has never been run so we run the sync
-    last_ext_group_sync = cc_pair.last_time_external_group_sync
-    if last_ext_group_sync is None:
-        return True
-
-    source_sync_period = EXTERNAL_GROUP_SYNC_PERIODS.get(cc_pair.connector.source)
-
-    # If EXTERNAL_GROUP_SYNC_PERIODS is None, we always run the sync.
-    if not source_sync_period:
-        return True
-
-    # If the last sync is greater than the full fetch period, we run the sync
-    next_sync = last_ext_group_sync + timedelta(seconds=source_sync_period)
-    if datetime.now(timezone.utc) >= next_sync:
-        return True
-
-    return False
-
-
-@shared_task(
-    name=DanswerCeleryTask.CHECK_FOR_EXTERNAL_GROUP_SYNC,
-    soft_time_limit=JOB_TIMEOUT,
-    bind=True,
-)
-def check_for_external_group_sync(self: Task, *, tenant_id: str | None) -> None:
-    r = get_redis_client(tenant_id=tenant_id)
-
-    lock_beat = r.lock(
-        DanswerRedisLocks.CHECK_CONNECTOR_EXTERNAL_GROUP_SYNC_BEAT_LOCK,
-        timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
-    )
-
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return
-
-        cc_pair_ids_to_sync: list[int] = []
-        with get_session_with_tenant(tenant_id) as db_session:
-            cc_pairs = get_all_auto_sync_cc_pairs(db_session)
-
-            # We only want to sync one cc_pair per source type in
-            # GROUP_PERMISSIONS_IS_CC_PAIR_AGNOSTIC
-            for source in GROUP_PERMISSIONS_IS_CC_PAIR_AGNOSTIC:
-                # These are ordered by cc_pair id so the first one is the one we want
-                cc_pairs_to_dedupe = get_cc_pairs_by_source(
-                    db_session, source, only_sync=True
-                )
-                # We only want to sync one cc_pair per source type
-                # in GROUP_PERMISSIONS_IS_CC_PAIR_AGNOSTIC so we dedupe here
-                for cc_pair_to_remove in cc_pairs_to_dedupe[1:]:
-                    cc_pairs = [
-                        cc_pair
-                        for cc_pair in cc_pairs
-                        if cc_pair.id != cc_pair_to_remove.id
-                    ]
-
-            for cc_pair in cc_pairs:
-                if _is_external_group_sync_due(cc_pair):
-                    cc_pair_ids_to_sync.append(cc_pair.id)
-
-        for cc_pair_id in cc_pair_ids_to_sync:
-            tasks_created = try_creating_external_group_sync_task(
-                self.app, cc_pair_id, r, tenant_id
-            )
-            if not tasks_created:
-                continue
-
-            task_logger.info(f"External group sync 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(f"Unexpected exception: tenant={tenant_id}")
-    finally:
-        if lock_beat.owned():
-            lock_beat.release()
-
-
-def try_creating_external_group_sync_task(
-    app: Celery,
-    cc_pair_id: int,
-    r: Redis,
-    tenant_id: str | None,
-) -> int | None:
-    """Returns an int if syncing is needed. The int represents the number of sync tasks generated.
-    Returns None if no syncing is required."""
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-
-    LOCK_TIMEOUT = 30
-
-    lock = r.lock(
-        DANSWER_REDIS_FUNCTION_LOCK_PREFIX + "try_generate_external_group_sync_tasks",
-        timeout=LOCK_TIMEOUT,
-    )
-
-    acquired = lock.acquire(blocking_timeout=LOCK_TIMEOUT / 2)
-    if not acquired:
-        return None
-
-    try:
-        # Dont kick off a new sync if the previous one is still running
-        if redis_connector.external_group_sync.fenced:
-            return None
-
-        redis_connector.external_group_sync.generator_clear()
-        redis_connector.external_group_sync.taskset_clear()
-
-        custom_task_id = f"{redis_connector.external_group_sync.taskset_key}_{uuid4()}"
-
-        result = app.send_task(
-            DanswerCeleryTask.CONNECTOR_EXTERNAL_GROUP_SYNC_GENERATOR_TASK,
-            kwargs=dict(
-                cc_pair_id=cc_pair_id,
-                tenant_id=tenant_id,
-            ),
-            queue=DanswerCeleryQueues.CONNECTOR_EXTERNAL_GROUP_SYNC,
-            task_id=custom_task_id,
-            priority=DanswerCeleryPriority.HIGH,
-        )
-
-        payload = RedisConnectorExternalGroupSyncPayload(
-            started=datetime.now(timezone.utc),
-            celery_task_id=result.id,
-        )
-
-        redis_connector.external_group_sync.set_fence(payload)
-
-    except Exception:
-        task_logger.exception(
-            f"Unexpected exception while trying to create external group sync task: cc_pair={cc_pair_id}"
-        )
-        return None
-    finally:
-        if lock.owned():
-            lock.release()
-
-    return 1
-
-
-@shared_task(
-    name=DanswerCeleryTask.CONNECTOR_EXTERNAL_GROUP_SYNC_GENERATOR_TASK,
-    acks_late=False,
-    soft_time_limit=JOB_TIMEOUT,
-    track_started=True,
-    trail=False,
-    bind=True,
-)
-def connector_external_group_sync_generator_task(
-    self: Task,
-    cc_pair_id: int,
-    tenant_id: str | None,
-) -> None:
-    """
-    Permission sync task that handles external group syncing for a given connector credential pair
-    This task assumes that the task has already been properly fenced
-    """
-
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-
-    r = get_redis_client(tenant_id=tenant_id)
-
-    lock: RedisLock = r.lock(
-        DanswerRedisLocks.CONNECTOR_EXTERNAL_GROUP_SYNC_LOCK_PREFIX
-        + f"_{redis_connector.id}",
-        timeout=CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT,
-    )
-
-    try:
-        acquired = lock.acquire(blocking=False)
-        if not acquired:
-            task_logger.warning(
-                f"External group sync task already running, exiting...: cc_pair={cc_pair_id}"
-            )
-            return None
-
-        with get_session_with_tenant(tenant_id) as db_session:
-            cc_pair = get_connector_credential_pair_from_id(cc_pair_id, db_session)
-            if cc_pair is None:
-                raise ValueError(
-                    f"No connector credential pair found for id: {cc_pair_id}"
-                )
-
-            source_type = cc_pair.connector.source
-
-            ext_group_sync_func = GROUP_PERMISSIONS_FUNC_MAP.get(source_type)
-            if ext_group_sync_func is None:
-                raise ValueError(
-                    f"No external group sync func found for {source_type} for cc_pair: {cc_pair_id}"
-                )
-
-            logger.info(
-                f"Syncing external groups for {source_type} for cc_pair: {cc_pair_id}"
-            )
-
-            external_user_groups: list[ExternalUserGroup] = ext_group_sync_func(cc_pair)
-
-            logger.info(
-                f"Syncing {len(external_user_groups)} external user groups for {source_type}"
-            )
-
-            replace_user__ext_group_for_cc_pair(
-                db_session=db_session,
-                cc_pair_id=cc_pair.id,
-                group_defs=external_user_groups,
-                source=cc_pair.connector.source,
-            )
-            logger.info(
-                f"Synced {len(external_user_groups)} external user groups for {source_type}"
-            )
-
-            mark_cc_pair_as_external_group_synced(db_session, cc_pair.id)
-    except Exception as e:
-        task_logger.exception(
-            f"Failed to run external group sync: cc_pair={cc_pair_id}"
-        )
-
-        redis_connector.external_group_sync.generator_clear()
-        redis_connector.external_group_sync.taskset_clear()
-        raise e
-    finally:
-        # we always want to clear the fence after the task is done or failed so it doesn't get stuck
-        redis_connector.external_group_sync.set_fence(None)
-        if lock.owned():
-            lock.release()

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

@@ -10,50 +10,41 @@ from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from redis import Redis
-from redis.exceptions import LockError
-from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 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
 from danswer.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX
 from danswer.configs.constants import DanswerCeleryPriority
 from danswer.configs.constants import DanswerCeleryQueues
-from danswer.configs.constants import DanswerCeleryTask
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.configs.constants import DocumentSource
-from danswer.db.connector import mark_ccpair_with_indexing_trigger
 from danswer.db.connector_credential_pair import fetch_connector_credential_pairs
 from danswer.db.connector_credential_pair import get_connector_credential_pair_from_id
 from danswer.db.engine import get_db_current_time
 from danswer.db.engine import get_session_with_tenant
 from danswer.db.enums import ConnectorCredentialPairStatus
-from danswer.db.enums import IndexingMode
 from danswer.db.enums import IndexingStatus
 from danswer.db.enums import IndexModelStatus
 from danswer.db.index_attempt import create_index_attempt
-from danswer.db.index_attempt import delete_index_attempt
-from danswer.db.index_attempt import get_all_index_attempts_by_status
 from danswer.db.index_attempt import get_index_attempt
 from danswer.db.index_attempt import get_last_attempt_for_cc_pair
-from danswer.db.index_attempt import mark_attempt_canceled
 from danswer.db.index_attempt import mark_attempt_failed
 from danswer.db.models import ConnectorCredentialPair
 from danswer.db.models import IndexAttempt
 from danswer.db.models import SearchSettings
-from danswer.db.search_settings import get_active_search_settings
 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.indexing.indexing_heartbeat import IndexingHeartbeatInterface
 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 RedisConnectorIndex
-from danswer.redis.redis_connector_index import RedisConnectorIndexPayload
+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
@@ -65,108 +56,41 @@ from shared_configs.configs import SENTRY_DSN
 logger = setup_logger()
 
 
-class IndexingCallback(IndexingHeartbeatInterface):
+class RunIndexingCallback(RunIndexingCallbackInterface):
     def __init__(
         self,
         stop_key: str,
         generator_progress_key: str,
-        redis_lock: RedisLock,
+        redis_lock: redis.lock.Lock,
         redis_client: Redis,
     ):
         super().__init__()
-        self.redis_lock: RedisLock = redis_lock
+        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
-        self.started: datetime = datetime.now(timezone.utc)
-        self.redis_lock.reacquire()
-
-        self.last_tag: str = "IndexingCallback.__init__"
-        self.last_lock_reacquire: datetime = datetime.now(timezone.utc)
 
     def should_stop(self) -> bool:
         if self.redis_client.exists(self.stop_key):
             return True
         return False
 
-    def progress(self, tag: str, amount: int) -> None:
-        try:
-            self.redis_lock.reacquire()
-            self.last_tag = tag
-            self.last_lock_reacquire = datetime.now(timezone.utc)
-        except LockError:
-            logger.exception(
-                f"IndexingCallback - lock.reacquire exceptioned. "
-                f"lock_timeout={self.redis_lock.timeout} "
-                f"start={self.started} "
-                f"last_tag={self.last_tag} "
-                f"last_reacquired={self.last_lock_reacquire} "
-                f"now={datetime.now(timezone.utc)}"
-            )
-            raise
-
+    def progress(self, amount: int) -> None:
+        self.redis_lock.reacquire()
         self.redis_client.incrby(self.generator_progress_key, amount)
 
 
-def get_unfenced_index_attempt_ids(db_session: Session, r: redis.Redis) -> list[int]:
-    """Gets a list of unfenced index attempts. Should not be possible, so we'd typically
-    want to clean them up.
-
-    Unfenced = attempt not in terminal state and fence does not exist.
-    """
-    unfenced_attempts: list[int] = []
-
-    # inner/outer/inner double check pattern to avoid race conditions when checking for
-    # bad state
-    # inner = index_attempt in non terminal state
-    # outer = r.fence_key down
-
-    # check the db for index attempts in a non terminal state
-    attempts: list[IndexAttempt] = []
-    attempts.extend(
-        get_all_index_attempts_by_status(IndexingStatus.NOT_STARTED, db_session)
-    )
-    attempts.extend(
-        get_all_index_attempts_by_status(IndexingStatus.IN_PROGRESS, db_session)
-    )
-
-    for attempt in attempts:
-        fence_key = RedisConnectorIndex.fence_key_with_ids(
-            attempt.connector_credential_pair_id, attempt.search_settings_id
-        )
-
-        # if the fence is down / doesn't exist, possible error but not confirmed
-        if r.exists(fence_key):
-            continue
-
-        # Between the time the attempts are first looked up and the time we see the fence down,
-        # the attempt may have completed and taken down the fence normally.
-
-        # We need to double check that the index attempt is still in a non terminal state
-        # and matches the original state, which confirms we are really in a bad state.
-        attempt_2 = get_index_attempt(db_session, attempt.id)
-        if not attempt_2:
-            continue
-
-        if attempt.status != attempt_2.status:
-            continue
-
-        unfenced_attempts.append(attempt.id)
-
-    return unfenced_attempts
-
-
 @shared_task(
-    name=DanswerCeleryTask.CHECK_FOR_INDEXING,
+    name="check_for_indexing",
     soft_time_limit=300,
     bind=True,
 )
 def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
     tasks_created = 0
-    locked = False
+
     r = get_redis_client(tenant_id=tenant_id)
 
-    lock_beat: RedisLock = r.lock(
+    lock_beat = r.lock(
         DanswerRedisLocks.CHECK_INDEXING_BEAT_LOCK,
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
@@ -176,9 +100,6 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
         if not lock_beat.acquire(blocking=False):
             return None
 
-        locked = True
-
-        # check for search settings swap
         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)
@@ -197,24 +118,26 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                         embedding_model=embedding_model,
                     )
 
-        # gather cc_pair_ids
         cc_pair_ids: list[int] = []
         with get_session_with_tenant(tenant_id) as db_session:
-            lock_beat.reacquire()
             cc_pairs = fetch_connector_credential_pairs(db_session)
             for cc_pair_entry in cc_pairs:
                 cc_pair_ids.append(cc_pair_entry.id)
 
-        # kick off index attempts
         for cc_pair_id in cc_pair_ids:
-            lock_beat.reacquire()
-
             redis_connector = RedisConnector(tenant_id, cc_pair_id)
             with get_session_with_tenant(tenant_id) as db_session:
-                search_settings_list: list[SearchSettings] = get_active_search_settings(
-                    db_session
-                )
-                for search_settings_instance in search_settings_list:
+                # 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
                     )
@@ -230,80 +153,33 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                     last_attempt = get_last_attempt_for_cc_pair(
                         cc_pair.id, search_settings_instance.id, db_session
                     )
-
-                    search_settings_primary = False
-                    if search_settings_instance.id == search_settings_list[0].id:
-                        search_settings_primary = True
-
                     if not _should_index(
                         cc_pair=cc_pair,
                         last_index=last_attempt,
                         search_settings_instance=search_settings_instance,
-                        search_settings_primary=search_settings_primary,
-                        secondary_index_building=len(search_settings_list) > 1,
+                        secondary_index_building=len(search_settings) > 1,
                         db_session=db_session,
                     ):
                         continue
 
-                    reindex = False
-                    if search_settings_instance.id == search_settings_list[0].id:
-                        # the indexing trigger is only checked and cleared with the primary search settings
-                        if cc_pair.indexing_trigger is not None:
-                            if cc_pair.indexing_trigger == IndexingMode.REINDEX:
-                                reindex = True
-
-                            task_logger.info(
-                                f"Connector indexing manual trigger detected: "
-                                f"cc_pair={cc_pair.id} "
-                                f"search_settings={search_settings_instance.id} "
-                                f"indexing_mode={cc_pair.indexing_trigger}"
-                            )
-
-                            mark_ccpair_with_indexing_trigger(
-                                cc_pair.id, None, db_session
-                            )
-
                     # 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,
-                        reindex,
+                        False,
                         db_session,
                         r,
                         tenant_id,
                     )
                     if attempt_id:
                         task_logger.info(
-                            f"Connector indexing queued: "
-                            f"index_attempt={attempt_id} "
+                            f"Indexing queued: index_attempt={attempt_id} "
                             f"cc_pair={cc_pair.id} "
-                            f"search_settings={search_settings_instance.id}"
+                            f"search_settings={search_settings_instance.id} "
                         )
                         tasks_created += 1
-
-        # Fail any index attempts in the DB that don't have fences
-        # This shouldn't ever happen!
-        with get_session_with_tenant(tenant_id) as db_session:
-            unfenced_attempt_ids = get_unfenced_index_attempt_ids(db_session, r)
-            for attempt_id in unfenced_attempt_ids:
-                lock_beat.reacquire()
-
-                attempt = get_index_attempt(db_session, attempt_id)
-                if not attempt:
-                    continue
-
-                failure_reason = (
-                    f"Unfenced index attempt found in DB: "
-                    f"index_attempt={attempt.id} "
-                    f"cc_pair={attempt.connector_credential_pair_id} "
-                    f"search_settings={attempt.search_settings_id}"
-                )
-                task_logger.error(failure_reason)
-                mark_attempt_failed(
-                    attempt.id, db_session, failure_reason=failure_reason
-                )
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
@@ -311,14 +187,8 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
     except Exception:
         task_logger.exception(f"Unexpected exception: tenant={tenant_id}")
     finally:
-        if locked:
-            if lock_beat.owned():
-                lock_beat.release()
-            else:
-                task_logger.error(
-                    "check_for_indexing - Lock not owned on completion: "
-                    f"tenant={tenant_id}"
-                )
+        if lock_beat.owned():
+            lock_beat.release()
 
     return tasks_created
 
@@ -327,7 +197,6 @@ def _should_index(
     cc_pair: ConnectorCredentialPair,
     last_index: IndexAttempt | None,
     search_settings_instance: SearchSettings,
-    search_settings_primary: bool,
     secondary_index_building: bool,
     db_session: Session,
 ) -> bool:
@@ -392,11 +261,6 @@ def _should_index(
     ):
         return False
 
-    if search_settings_primary:
-        if cc_pair.indexing_trigger is not None:
-            # if a manual indexing trigger is on the cc pair, honor it for primary search settings
-            return True
-
     # if no attempt has ever occurred, we should index regardless of refresh_freq
     if not last_index:
         return True
@@ -429,11 +293,10 @@ def try_creating_indexing_task(
     """
 
     LOCK_TIMEOUT = 30
-    index_attempt_id: int | None = None
 
     # we need to serialize any attempt to trigger indexing since it can be triggered
     # either via celery beat or manually (API call)
-    lock: RedisLock = r.lock(
+    lock = r.lock(
         DANSWER_REDIS_FUNCTION_LOCK_PREFIX + "try_creating_indexing_task",
         timeout=LOCK_TIMEOUT,
     )
@@ -462,7 +325,7 @@ def try_creating_indexing_task(
         redis_connector_index.generator_clear()
 
         # set a basic fence to start
-        payload = RedisConnectorIndexPayload(
+        payload = RedisConnectorIndexingFenceData(
             index_attempt_id=None,
             started=None,
             submitted=datetime.now(timezone.utc),
@@ -484,10 +347,8 @@ def try_creating_indexing_task(
 
         custom_task_id = redis_connector_index.generate_generator_task_id()
 
-        # when the task is sent, we have yet to finish setting up the fence
-        # therefore, the task must contain code that blocks until the fence is ready
         result = celery_app.send_task(
-            DanswerCeleryTask.CONNECTOR_INDEXING_PROXY_TASK,
+            "connector_indexing_proxy_task",
             kwargs=dict(
                 index_attempt_id=index_attempt_id,
                 cc_pair_id=cc_pair.id,
@@ -505,17 +366,15 @@ def try_creating_indexing_task(
         payload.index_attempt_id = index_attempt_id
         payload.celery_task_id = result.id
         redis_connector_index.set_fence(payload)
+
     except Exception:
+        redis_connector_index.set_fence(payload)
         task_logger.exception(
-            f"try_creating_indexing_task - Unexpected exception: "
+            f"Unexpected exception: "
             f"tenant={tenant_id} "
             f"cc_pair={cc_pair.id} "
             f"search_settings={search_settings.id}"
         )
-
-        if index_attempt_id is not None:
-            delete_index_attempt(db_session, index_attempt_id)
-        redis_connector_index.set_fence(None)
         return None
     finally:
         if lock.owned():
@@ -524,14 +383,8 @@ def try_creating_indexing_task(
     return index_attempt_id
 
 
-@shared_task(
-    name=DanswerCeleryTask.CONNECTOR_INDEXING_PROXY_TASK,
-    bind=True,
-    acks_late=False,
-    track_started=True,
-)
+@shared_task(name="connector_indexing_proxy_task", acks_late=False, track_started=True)
 def connector_indexing_proxy_task(
-    self: Task,
     index_attempt_id: int,
     cc_pair_id: int,
     search_settings_id: int,
@@ -539,19 +392,15 @@ def connector_indexing_proxy_task(
 ) -> None:
     """celery tasks are forked, but forking is unstable.  This proxies work to a spawned task."""
     task_logger.info(
-        f"Indexing watchdog - starting: attempt={index_attempt_id} "
+        f"Indexing proxy - starting: attempt={index_attempt_id} "
         f"tenant={tenant_id} "
         f"cc_pair={cc_pair_id} "
         f"search_settings={search_settings_id}"
     )
-
-    if not self.request.id:
-        task_logger.error("self.request.id is None!")
-
     client = SimpleJobClient()
 
     job = client.submit(
-        connector_indexing_task_wrapper,
+        connector_indexing_task,
         index_attempt_id,
         cc_pair_id,
         search_settings_id,
@@ -562,7 +411,7 @@ def connector_indexing_proxy_task(
 
     if not job:
         task_logger.info(
-            f"Indexing watchdog - spawn failed: attempt={index_attempt_id} "
+            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}"
@@ -570,78 +419,31 @@ def connector_indexing_proxy_task(
         return
 
     task_logger.info(
-        f"Indexing watchdog - spawn succeeded: attempt={index_attempt_id} "
+        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}"
     )
 
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-    redis_connector_index = redis_connector.new_index(search_settings_id)
-
     while True:
-        sleep(5)
-
-        if self.request.id and redis_connector_index.terminating(self.request.id):
-            task_logger.warning(
-                "Indexing watchdog - termination signal detected: "
-                f"attempt={index_attempt_id} "
-                f"tenant={tenant_id} "
-                f"cc_pair={cc_pair_id} "
-                f"search_settings={search_settings_id}"
-            )
-
-            try:
-                with get_session_with_tenant(tenant_id) as db_session:
-                    mark_attempt_canceled(
-                        index_attempt_id,
-                        db_session,
-                        "Connector termination signal detected",
-                    )
-            finally:
-                # if the DB exceptions, we'll just get an unfriendly failure message
-                # in the UI instead of the cancellation message
-                logger.exception(
-                    "Indexing watchdog - transient exception marking index attempt as canceled: "
-                    f"attempt={index_attempt_id} "
-                    f"tenant={tenant_id} "
-                    f"cc_pair={cc_pair_id} "
-                    f"search_settings={search_settings_id}"
-                )
-
-                job.cancel()
-
-            break
+        sleep(10)
 
+        # do nothing for ongoing jobs that haven't been stopped
         if not job.done():
-            # if the spawned task is still running, restart the check once again
-            # if the index attempt is not in a finished status
-            try:
-                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
-            except Exception:
-                # if the DB exceptioned, just restart the check.
-                # polling the index attempt status doesn't need to be strongly consistent
-                logger.exception(
-                    "Indexing watchdog - transient exception looking up index attempt: "
-                    f"attempt={index_attempt_id} "
-                    f"tenant={tenant_id} "
-                    f"cc_pair={cc_pair_id} "
-                    f"search_settings={search_settings_id}"
+            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
                 )
-                continue
+
+                if not index_attempt:
+                    continue
+
+                if not index_attempt.is_finished():
+                    continue
 
         if job.status == "error":
             task_logger.error(
-                "Indexing watchdog - spawned task exceptioned: "
+                f"Indexing proxy - spawned task exceptioned: "
                 f"attempt={index_attempt_id} "
                 f"tenant={tenant_id} "
                 f"cc_pair={cc_pair_id} "
@@ -653,7 +455,7 @@ def connector_indexing_proxy_task(
         break
 
     task_logger.info(
-        f"Indexing watchdog - finished: attempt={index_attempt_id} "
+        f"Indexing proxy - finished: attempt={index_attempt_id} "
         f"tenant={tenant_id} "
         f"cc_pair={cc_pair_id} "
         f"search_settings={search_settings_id}"
@@ -661,38 +463,6 @@ def connector_indexing_proxy_task(
     return
 
 
-def connector_indexing_task_wrapper(
-    index_attempt_id: int,
-    cc_pair_id: int,
-    search_settings_id: int,
-    tenant_id: str | None,
-    is_ee: bool,
-) -> int | None:
-    """Just wraps connector_indexing_task so we can log any exceptions before
-    re-raising it."""
-    result: int | None = None
-
-    try:
-        result = connector_indexing_task(
-            index_attempt_id,
-            cc_pair_id,
-            search_settings_id,
-            tenant_id,
-            is_ee,
-        )
-    except:
-        logger.exception(
-            f"connector_indexing_task exceptioned: "
-            f"tenant={tenant_id} "
-            f"index_attempt={index_attempt_id} "
-            f"cc_pair={cc_pair_id} "
-            f"search_settings={search_settings_id}"
-        )
-        raise
-
-    return result
-
-
 def connector_indexing_task(
     index_attempt_id: int,
     cc_pair_id: int,
@@ -729,8 +499,7 @@ def connector_indexing_task(
         logger.debug("Sentry DSN not provided, skipping Sentry initialization")
 
     logger.info(
-        f"Indexing spawned task starting: "
-        f"attempt={index_attempt_id} "
+        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}"
@@ -747,7 +516,6 @@ def connector_indexing_task(
     if redis_connector.delete.fenced:
         raise RuntimeError(
             f"Indexing will not start because connector deletion is in progress: "
-            f"attempt={index_attempt_id} "
             f"cc_pair={cc_pair_id} "
             f"fence={redis_connector.delete.fence_key}"
         )
@@ -755,18 +523,18 @@ def connector_indexing_task(
     if redis_connector.stop.fenced:
         raise RuntimeError(
             f"Indexing will not start because a connector stop signal was detected: "
-            f"attempt={index_attempt_id} "
             f"cc_pair={cc_pair_id} "
             f"fence={redis_connector.stop.fence_key}"
         )
 
     while True:
-        if not redis_connector_index.fenced:  # The fence must exist
+        # 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  # The payload must exist
+        payload = redis_connector_index.payload
         if not payload:
             raise ValueError("connector_indexing_task: payload invalid or not found")
 
@@ -789,19 +557,16 @@ def connector_indexing_task(
         )
         break
 
-    # set thread_local=False since we don't control what thread the indexing/pruning
-    # might run our callback with
-    lock: RedisLock = r.lock(
+    lock = r.lock(
         redis_connector_index.generator_lock_key,
         timeout=CELERY_INDEXING_LOCK_TIMEOUT,
-        thread_local=False,
     )
 
     acquired = lock.acquire(blocking=False)
     if not acquired:
         logger.warning(
             f"Indexing task already running, exiting...: "
-            f"index_attempt={index_attempt_id} cc_pair={cc_pair_id} search_settings={search_settings_id}"
+            f"cc_pair={cc_pair_id} search_settings={search_settings_id}"
         )
         return None
 
@@ -836,7 +601,7 @@ def connector_indexing_task(
                 )
 
         # define a callback class
-        callback = IndexingCallback(
+        callback = RunIndexingCallback(
             redis_connector.stop.fence_key,
             redis_connector_index.generator_progress_key,
             lock,

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

@@ -13,13 +13,12 @@ from sqlalchemy.orm import Session
 
 from danswer.background.celery.apps.app_base import task_logger
 from danswer.configs.app_configs import JOB_TIMEOUT
-from danswer.configs.constants import DanswerCeleryTask
 from danswer.configs.constants import PostgresAdvisoryLocks
 from danswer.db.engine import get_session_with_tenant
 
 
 @shared_task(
-    name=DanswerCeleryTask.KOMBU_MESSAGE_CLEANUP_TASK,
+    name="kombu_message_cleanup_task",
     soft_time_limit=JOB_TIMEOUT,
     bind=True,
     base=AbortableTask,

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

@@ -8,12 +8,11 @@ from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from redis import Redis
-from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 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 IndexingCallback
+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
@@ -21,7 +20,6 @@ from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
 from danswer.configs.constants import DANSWER_REDIS_FUNCTION_LOCK_PREFIX
 from danswer.configs.constants import DanswerCeleryPriority
 from danswer.configs.constants import DanswerCeleryQueues
-from danswer.configs.constants import DanswerCeleryTask
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.connectors.factory import instantiate_connector
 from danswer.connectors.models import InputType
@@ -40,44 +38,8 @@ from danswer.utils.logger import setup_logger
 logger = setup_logger()
 
 
-def _is_pruning_due(cc_pair: ConnectorCredentialPair) -> bool:
-    """Returns boolean indicating if pruning is due.
-
-    Next pruning time is calculated as a delta from the last successful prune, or the
-    last successful indexing if pruning has never succeeded.
-
-    TODO(rkuo): consider whether we should allow pruning to be immediately rescheduled
-    if pruning fails (which is what it does now). A backoff could be reasonable.
-    """
-
-    # skip pruning if no prune frequency is set
-    # pruning can still be forced via the API which will run a pruning task directly
-    if not cc_pair.connector.prune_freq:
-        return False
-
-    # skip pruning if not active
-    if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:
-        return False
-
-    # skip pruning if the next scheduled prune time hasn't been reached yet
-    last_pruned = cc_pair.last_pruned
-    if not last_pruned:
-        if not cc_pair.last_successful_index_time:
-            # if we've never indexed, we can't prune
-            return False
-
-        # if never pruned, use the last time the connector indexed successfully
-        last_pruned = cc_pair.last_successful_index_time
-
-    next_prune = last_pruned + timedelta(seconds=cc_pair.connector.prune_freq)
-    if datetime.now(timezone.utc) < next_prune:
-        return False
-
-    return True
-
-
 @shared_task(
-    name=DanswerCeleryTask.CHECK_FOR_PRUNING,
+    name="check_for_pruning",
     soft_time_limit=JOB_TIMEOUT,
     bind=True,
 )
@@ -107,7 +69,7 @@ def check_for_pruning(self: Task, *, tenant_id: str | None) -> None:
                 if not cc_pair:
                     continue
 
-                if not _is_pruning_due(cc_pair):
+                if not is_pruning_due(cc_pair, db_session, r):
                     continue
 
                 tasks_created = try_creating_prune_generator_task(
@@ -128,6 +90,47 @@ def check_for_pruning(self: Task, *, tenant_id: str | None) -> None:
             lock_beat.release()
 
 
+def is_pruning_due(
+    cc_pair: ConnectorCredentialPair,
+    db_session: Session,
+    r: Redis,
+) -> bool:
+    """Returns an int if pruning is triggered.
+    The int represents the number of prune tasks generated (in this case, only one
+    because the task is a long running generator task.)
+    Returns None if no pruning is triggered (due to not being needed or
+    other reasons such as simultaneous pruning restrictions.
+
+    Checks for scheduling related conditions, then delegates the rest of the checks to
+    try_creating_prune_generator_task.
+    """
+
+    # skip pruning if no prune frequency is set
+    # pruning can still be forced via the API which will run a pruning task directly
+    if not cc_pair.connector.prune_freq:
+        return False
+
+    # skip pruning if not active
+    if cc_pair.status != ConnectorCredentialPairStatus.ACTIVE:
+        return False
+
+    # skip pruning if the next scheduled prune time hasn't been reached yet
+    last_pruned = cc_pair.last_pruned
+    if not last_pruned:
+        if not cc_pair.last_successful_index_time:
+            # if we've never indexed, we can't prune
+            return False
+
+        # if never pruned, use the last time the connector indexed successfully
+        last_pruned = cc_pair.last_successful_index_time
+
+    next_prune = last_pruned + timedelta(seconds=cc_pair.connector.prune_freq)
+    if datetime.now(timezone.utc) < next_prune:
+        return False
+
+    return True
+
+
 def try_creating_prune_generator_task(
     celery_app: Celery,
     cc_pair: ConnectorCredentialPair,
@@ -163,16 +166,10 @@ def try_creating_prune_generator_task(
         return None
 
     try:
-        # skip pruning if already pruning
-        if redis_connector.prune.fenced:
+        if redis_connector.prune.fenced:  # skip pruning if already pruning
             return None
 
-        # skip pruning if the cc_pair is deleting
-        if redis_connector.delete.fenced:
-            return None
-
-        # skip pruning if doc permissions sync is running
-        if redis_connector.permissions.fenced:
+        if redis_connector.delete.fenced:  # skip pruning if the cc_pair is deleting
             return None
 
         db_session.refresh(cc_pair)
@@ -186,7 +183,7 @@ def try_creating_prune_generator_task(
         custom_task_id = f"{redis_connector.prune.generator_task_key}_{uuid4()}"
 
         celery_app.send_task(
-            DanswerCeleryTask.CONNECTOR_PRUNING_GENERATOR_TASK,
+            "connector_pruning_generator_task",
             kwargs=dict(
                 cc_pair_id=cc_pair.id,
                 connector_id=cc_pair.connector_id,
@@ -211,7 +208,7 @@ def try_creating_prune_generator_task(
 
 
 @shared_task(
-    name=DanswerCeleryTask.CONNECTOR_PRUNING_GENERATOR_TASK,
+    name="connector_pruning_generator_task",
     acks_late=False,
     soft_time_limit=JOB_TIMEOUT,
     track_started=True,
@@ -234,18 +231,13 @@ def connector_pruning_generator_task(
     pruning_ctx_dict["request_id"] = self.request.id
     pruning_ctx.set(pruning_ctx_dict)
 
-    task_logger.info(f"Pruning generator starting: cc_pair={cc_pair_id}")
-
     redis_connector = RedisConnector(tenant_id, cc_pair_id)
 
     r = get_redis_client(tenant_id=tenant_id)
 
-    # set thread_local=False since we don't control what thread the indexing/pruning
-    # might run our callback with
-    lock: RedisLock = r.lock(
+    lock = r.lock(
         DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{redis_connector.id}",
         timeout=CELERY_PRUNING_LOCK_TIMEOUT,
-        thread_local=False,
     )
 
     acquired = lock.acquire(blocking=False)
@@ -269,11 +261,6 @@ def connector_pruning_generator_task(
                 )
                 return
 
-            task_logger.info(
-                f"Pruning generator running connector: "
-                f"cc_pair={cc_pair_id} "
-                f"connector_source={cc_pair.connector.source}"
-            )
             runnable_connector = instantiate_connector(
                 db_session,
                 cc_pair.connector.source,
@@ -282,13 +269,12 @@ def connector_pruning_generator_task(
                 cc_pair.credential,
             )
 
-            callback = IndexingCallback(
+            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, callback
@@ -310,8 +296,8 @@ def connector_pruning_generator_task(
             task_logger.info(
                 f"Pruning set collected: "
                 f"cc_pair={cc_pair_id} "
-                f"connector_source={cc_pair.connector.source} "
-                f"docs_to_remove={len(doc_ids_to_remove)}"
+                f"docs_to_remove={len(doc_ids_to_remove)} "
+                f"doc_source={cc_pair.connector.source}"
             )
 
             task_logger.info(
@@ -334,10 +320,10 @@ def connector_pruning_generator_task(
             f"Failed to run pruning: cc_pair={cc_pair_id} connector={connector_id}"
         )
 
-        redis_connector.prune.reset()
+        redis_connector.prune.generator_clear()
+        redis_connector.prune.taskset_clear()
+        redis_connector.prune.set_fence(False)
         raise e
     finally:
         if lock.owned():
             lock.release()
-
-        task_logger.info(f"Pruning generator finished: cc_pair={cc_pair_id}")

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

@@ -9,7 +9,6 @@ from tenacity import RetryError
 from danswer.access.access import get_access_for_document
 from danswer.background.celery.apps.app_base import task_logger
 from danswer.background.celery.tasks.shared.RetryDocumentIndex import RetryDocumentIndex
-from danswer.configs.constants import DanswerCeleryTask
 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
@@ -32,7 +31,7 @@ LIGHT_TIME_LIMIT = LIGHT_SOFT_TIME_LIMIT + 15
 
 
 @shared_task(
-    name=DanswerCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_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,
@@ -60,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.debug(f"Task start: tenant={tenant_id} doc={document_id}")
+    task_logger.info(f"tenant={tenant_id} doc={document_id}")
 
     try:
         with get_session_with_tenant(tenant_id) as db_session:
@@ -142,9 +141,7 @@ def document_by_cc_pair_cleanup_task(
         return False
     except Exception as ex:
         if isinstance(ex, RetryError):
-            task_logger.warning(
-                f"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}"
-            )
+            task_logger.info(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()
@@ -174,21 +171,11 @@ def document_by_cc_pair_cleanup_task(
         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.warning(
-                f"Max celery task retries reached. Marking doc as dirty for 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) as db_session:
-                # delete the cc pair relationship now and let reconciliation clean it up
-                # in vespa
-                delete_document_by_connector_credential_pair__no_commit(
-                    db_session=db_session,
-                    document_id=document_id,
-                    connector_credential_pair_identifier=ConnectorCredentialPairIdentifier(
-                        connector_id=connector_id,
-                        credential_id=credential_id,
-                    ),
-                )
+            with get_session_with_tenant(tenant_id):
                 mark_document_as_modified(document_id, db_session)
         return False
 

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

@@ -5,6 +5,7 @@ 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
@@ -12,7 +13,6 @@ from celery.exceptions import SoftTimeLimitExceeded
 from celery.result import AsyncResult
 from celery.states import READY_STATES
 from redis import Redis
-from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 from tenacity import RetryError
 
@@ -25,10 +25,8 @@ 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
-from danswer.configs.constants import DanswerCeleryTask
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.db.connector import fetch_connector_by_id
-from danswer.db.connector import mark_cc_pair_as_permissions_synced
 from danswer.db.connector import mark_ccpair_as_pruned
 from danswer.db.connector_credential_pair import add_deletion_failure_message
 from danswer.db.connector_credential_pair import (
@@ -49,19 +47,17 @@ from danswer.db.document_set import mark_document_set_as_synced
 from danswer.db.engine import get_session_with_tenant
 from danswer.db.enums import IndexingStatus
 from danswer.db.index_attempt import delete_index_attempts
+from danswer.db.index_attempt import get_all_index_attempts_by_status
 from danswer.db.index_attempt import get_index_attempt
 from danswer.db.index_attempt import mark_attempt_failed
 from danswer.db.models import DocumentSet
+from danswer.db.models import IndexAttempt
 from danswer.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_doc_perm_sync import RedisConnectorPermissionSync
-from danswer.redis.redis_connector_doc_perm_sync import (
-    RedisConnectorPermissionSyncPayload,
-)
 from danswer.redis.redis_connector_index import RedisConnectorIndex
 from danswer.redis.redis_connector_prune import RedisConnectorPrune
 from danswer.redis.redis_document_set import RedisDocumentSet
@@ -81,7 +77,7 @@ logger = setup_logger()
 # celery auto associates tasks created inside another task,
 # which bloats the result metadata considerably. trail=False prevents this.
 @shared_task(
-    name=DanswerCeleryTask.CHECK_FOR_VESPA_SYNC_TASK,
+    name="check_for_vespa_sync_task",
     soft_time_limit=JOB_TIMEOUT,
     trail=False,
     bind=True,
@@ -166,7 +162,7 @@ def try_generate_stale_document_sync_tasks(
     celery_app: Celery,
     db_session: Session,
     r: Redis,
-    lock_beat: RedisLock,
+    lock_beat: redis.lock.Lock,
     tenant_id: str | None,
 ) -> int | None:
     # the fence is up, do nothing
@@ -184,12 +180,7 @@ def try_generate_stale_document_sync_tasks(
         f"Stale documents found (at least {stale_doc_count}). Generating sync tasks by cc pair."
     )
 
-    task_logger.info(
-        "RedisConnector.generate_tasks starting by cc_pair. "
-        "Documents spanning multiple cc_pairs will only be synced once."
-    )
-
-    docs_to_skip: set[str] = set()
+    task_logger.info("RedisConnector.generate_tasks starting by cc_pair.")
 
     # rkuo: we could technically sync all stale docs in one big pass.
     # but I feel it's more understandable to group the docs by cc_pair
@@ -197,21 +188,22 @@ def try_generate_stale_document_sync_tasks(
     cc_pairs = get_connector_credential_pairs(db_session)
     for cc_pair in cc_pairs:
         rc = RedisConnectorCredentialPair(tenant_id, cc_pair.id)
-        rc.set_skip_docs(docs_to_skip)
-        result = rc.generate_tasks(celery_app, db_session, r, lock_beat, tenant_id)
+        tasks_generated = rc.generate_tasks(
+            celery_app, db_session, r, lock_beat, tenant_id
+        )
 
-        if result is None:
+        if tasks_generated is None:
             continue
 
-        if result[1] == 0:
+        if tasks_generated == 0:
             continue
 
         task_logger.info(
             f"RedisConnector.generate_tasks finished for single cc_pair. "
-            f"cc_pair={cc_pair.id} tasks_generated={result[0]} tasks_possible={result[1]}"
+            f"cc_pair_id={cc_pair.id} tasks_generated={tasks_generated}"
         )
 
-        total_tasks_generated += result[0]
+        total_tasks_generated += tasks_generated
 
     task_logger.info(
         f"RedisConnector.generate_tasks finished for all cc_pairs. total_tasks_generated={total_tasks_generated}"
@@ -226,7 +218,7 @@ def try_generate_document_set_sync_tasks(
     document_set_id: int,
     db_session: Session,
     r: Redis,
-    lock_beat: RedisLock,
+    lock_beat: redis.lock.Lock,
     tenant_id: str | None,
 ) -> int | None:
     lock_beat.reacquire()
@@ -254,11 +246,12 @@ def try_generate_document_set_sync_tasks(
     )
 
     # Add all documents that need to be updated into the queue
-    result = rds.generate_tasks(celery_app, db_session, r, lock_beat, tenant_id)
-    if result is None:
+    tasks_generated = rds.generate_tasks(
+        celery_app, db_session, r, lock_beat, tenant_id
+    )
+    if tasks_generated is None:
         return None
 
-    tasks_generated = result[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.
@@ -267,7 +260,7 @@ def try_generate_document_set_sync_tasks(
 
     task_logger.info(
         f"RedisDocumentSet.generate_tasks finished. "
-        f"document_set={document_set.id} tasks_generated={tasks_generated}"
+        f"document_set_id={document_set.id} tasks_generated={tasks_generated}"
     )
 
     # set this only after all tasks have been added
@@ -280,7 +273,7 @@ def try_generate_user_group_sync_tasks(
     usergroup_id: int,
     db_session: Session,
     r: Redis,
-    lock_beat: RedisLock,
+    lock_beat: redis.lock.Lock,
     tenant_id: str | None,
 ) -> int | None:
     lock_beat.reacquire()
@@ -309,11 +302,12 @@ def try_generate_user_group_sync_tasks(
     task_logger.info(
         f"RedisUserGroup.generate_tasks starting. usergroup_id={usergroup.id}"
     )
-    result = rug.generate_tasks(celery_app, db_session, r, lock_beat, tenant_id)
-    if result is None:
+    tasks_generated = rug.generate_tasks(
+        celery_app, db_session, r, lock_beat, tenant_id
+    )
+    if tasks_generated is None:
         return None
 
-    tasks_generated = result[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.
@@ -322,7 +316,7 @@ def try_generate_user_group_sync_tasks(
 
     task_logger.info(
         f"RedisUserGroup.generate_tasks finished. "
-        f"usergroup={usergroup.id} tasks_generated={tasks_generated}"
+        f"usergroup_id={usergroup.id} tasks_generated={tasks_generated}"
     )
 
     # set this only after all tasks have been added
@@ -442,22 +436,11 @@ def monitor_connector_deletion_taskset(
                 db_session, cc_pair.connector_id, cc_pair.credential_id
             )
             if len(doc_ids) > 0:
-                # NOTE(rkuo): if this happens, documents somehow got added while
-                # deletion was in progress. Likely a bug gating off pruning and indexing
-                # work before deletion starts.
+                # if this happens, documents somehow got added while deletion was in progress. Likely a bug
+                # gating off pruning and indexing work before deletion starts
                 task_logger.warning(
-                    "Connector deletion - documents still found after taskset completion. "
-                    "Clearing the current deletion attempt and allowing deletion to restart: "
-                    f"cc_pair={cc_pair_id} "
-                    f"docs_deleted={fence_data.num_tasks} "
-                    f"docs_remaining={len(doc_ids)}"
-                )
-
-                # We don't want to waive off why we get into this state, but resetting
-                # our attempt and letting the deletion restart is a good way to recover
-                redis_connector.delete.reset()
-                raise RuntimeError(
-                    "Connector deletion - documents still found after taskset completion"
+                    f"Connector deletion - documents still found after taskset completion: "
+                    f"cc_pair={cc_pair_id} num={len(doc_ids)}"
                 )
 
             # clean up the rest of the related Postgres entities
@@ -521,7 +504,8 @@ def monitor_connector_deletion_taskset(
         f"docs_deleted={fence_data.num_tasks}"
     )
 
-    redis_connector.delete.reset()
+    redis_connector.delete.taskset_clear()
+    redis_connector.delete.set_fence(None)
 
 
 def monitor_ccpair_pruning_taskset(
@@ -562,45 +546,6 @@ def monitor_ccpair_pruning_taskset(
     redis_connector.prune.set_fence(False)
 
 
-def monitor_ccpair_permissions_taskset(
-    tenant_id: str | None, key_bytes: bytes, r: Redis, db_session: Session
-) -> None:
-    fence_key = key_bytes.decode("utf-8")
-    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_permissions_taskset: could not parse cc_pair_id from {fence_key}"
-        )
-        return
-
-    cc_pair_id = int(cc_pair_id_str)
-
-    redis_connector = RedisConnector(tenant_id, cc_pair_id)
-    if not redis_connector.permissions.fenced:
-        return
-
-    initial = redis_connector.permissions.generator_complete
-    if initial is None:
-        return
-
-    remaining = redis_connector.permissions.get_remaining()
-    task_logger.info(
-        f"Permissions sync progress: cc_pair={cc_pair_id} remaining={remaining} initial={initial}"
-    )
-    if remaining > 0:
-        return
-
-    payload: RedisConnectorPermissionSyncPayload | None = (
-        redis_connector.permissions.payload
-    )
-    start_time: datetime | None = payload.started if payload else None
-
-    mark_cc_pair_as_permissions_synced(db_session, int(cc_pair_id), start_time)
-    task_logger.info(f"Successfully synced permissions for cc_pair={cc_pair_id}")
-
-    redis_connector.permissions.reset()
-
-
 def monitor_ccpair_indexing_taskset(
     tenant_id: str | None, key_bytes: bytes, r: Redis, db_session: Session
 ) -> None:
@@ -635,8 +580,8 @@ def monitor_ccpair_indexing_taskset(
     progress = redis_connector_index.get_progress()
     if progress is not None:
         task_logger.info(
-            f"Connector indexing progress: cc_pair={cc_pair_id} "
-            f"search_settings={search_settings_id} "
+            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}"
         )
@@ -645,62 +590,39 @@ def monitor_ccpair_indexing_taskset(
         # the task is still setting up
         return
 
+    # Read result state BEFORE generator_complete_key to avoid a race condition
     # never use any blocking methods on the result from inside a task!
     result: AsyncResult = AsyncResult(payload.celery_task_id)
+    result_state = result.state
 
-    # inner/outer/inner double check pattern to avoid race conditions when checking for
-    # bad state
-
-    # inner = get_completion / generator_complete not signaled
-    # outer = result.state in READY state
     status_int = redis_connector_index.get_completion()
-    if status_int is None:  # inner signal not set ... possible error
-        task_state = result.state
-        if (
-            task_state in READY_STATES
-        ):  # outer signal in terminal state ... possible error
-            # Now double check!
-            if redis_connector_index.get_completion() is None:
-                # inner signal still not set (and cannot change when outer result_state is READY)
-                # Task is finished but generator complete isn't set.
-                # We have a problem! Worker may have crashed.
-                task_result = str(result.result)
-                task_traceback = str(result.traceback)
+    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
+            task_logger.info(
+                f"Connector indexing aborted: "
+                f"cc_pair_id={cc_pair_id} "
+                f"search_settings_id={search_settings_id} "
+                f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
+            )
 
-                msg = (
-                    f"Connector indexing aborted or exceptioned: "
-                    f"attempt={payload.index_attempt_id} "
-                    f"celery_task={payload.celery_task_id} "
-                    f"cc_pair={cc_pair_id} "
-                    f"search_settings={search_settings_id} "
-                    f"elapsed_submitted={elapsed_submitted.total_seconds():.2f} "
-                    f"result.state={task_state} "
-                    f"result.result={task_result} "
-                    f"result.traceback={task_traceback}"
+            index_attempt = get_index_attempt(db_session, payload.index_attempt_id)
+            if index_attempt:
+                mark_attempt_failed(
+                    index_attempt_id=payload.index_attempt_id,
+                    db_session=db_session,
+                    failure_reason="Connector indexing aborted or exceptioned.",
                 )
-                task_logger.warning(msg)
 
-                index_attempt = get_index_attempt(db_session, payload.index_attempt_id)
-                if index_attempt:
-                    if (
-                        index_attempt.status != IndexingStatus.CANCELED
-                        and index_attempt.status != IndexingStatus.FAILED
-                    ):
-                        mark_attempt_failed(
-                            index_attempt_id=payload.index_attempt_id,
-                            db_session=db_session,
-                            failure_reason=msg,
-                        )
-
-                redis_connector_index.reset()
+            redis_connector_index.reset()
         return
 
     status_enum = HTTPStatus(status_int)
 
     task_logger.info(
-        f"Connector indexing finished: cc_pair={cc_pair_id} "
-        f"search_settings={search_settings_id} "
-        f"progress={progress} "
+        f"Connector indexing finished: cc_pair_id={cc_pair_id} "
+        f"search_settings_id={search_settings_id} "
         f"status={status_enum.name} "
         f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
     )
@@ -708,7 +630,7 @@ def monitor_ccpair_indexing_taskset(
     redis_connector_index.reset()
 
 
-@shared_task(name=DanswerCeleryTask.MONITOR_VESPA_SYNC, soft_time_limit=300, bind=True)
+@shared_task(name="monitor_vespa_sync", soft_time_limit=300, bind=True)
 def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     """This is a celery beat task that monitors and finalizes metadata sync tasksets.
     It scans for fence values and then gets the counts of any associated tasksets.
@@ -721,7 +643,7 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     """
     r = get_redis_client(tenant_id=tenant_id)
 
-    lock_beat: RedisLock = r.lock(
+    lock_beat: redis.lock.Lock = r.lock(
         DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK,
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
@@ -733,7 +655,7 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
 
         # print current queue lengths
         r_celery = self.app.broker_connection().channel().client  # type: ignore
-        n_celery = celery_get_queue_length("celery", r_celery)
+        n_celery = celery_get_queue_length("celery", r)
         n_indexing = celery_get_queue_length(
             DanswerCeleryQueues.CONNECTOR_INDEXING, r_celery
         )
@@ -746,19 +668,41 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
         n_pruning = celery_get_queue_length(
             DanswerCeleryQueues.CONNECTOR_PRUNING, r_celery
         )
-        n_permissions_sync = celery_get_queue_length(
-            DanswerCeleryQueues.CONNECTOR_DOC_PERMISSIONS_SYNC, r_celery
-        )
 
         task_logger.info(
             f"Queue lengths: celery={n_celery} "
             f"indexing={n_indexing} "
             f"sync={n_sync} "
             f"deletion={n_deletion} "
-            f"pruning={n_pruning} "
-            f"permissions_sync={n_permissions_sync} "
+            f"pruning={n_pruning}"
         )
 
+        # do some cleanup before clearing fences
+        # check the db for any outstanding index attempts
+        with get_session_with_tenant(tenant_id) as db_session:
+            attempts: list[IndexAttempt] = []
+            attempts.extend(
+                get_all_index_attempts_by_status(IndexingStatus.NOT_STARTED, db_session)
+            )
+            attempts.extend(
+                get_all_index_attempts_by_status(IndexingStatus.IN_PROGRESS, db_session)
+            )
+
+            for a in attempts:
+                # if attempts exist in the db but we don't detect them in redis, mark them as failed
+                fence_key = RedisConnectorIndex.fence_key_with_ids(
+                    a.connector_credential_pair_id, a.search_settings_id
+                )
+                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)
@@ -797,12 +741,6 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
             with get_session_with_tenant(tenant_id) as db_session:
                 monitor_ccpair_indexing_taskset(tenant_id, key_bytes, r, db_session)
 
-        lock_beat.reacquire()
-        for key_bytes in r.scan_iter(RedisConnectorPermissionSync.FENCE_PREFIX + "*"):
-            lock_beat.reacquire()
-            with get_session_with_tenant(tenant_id) as db_session:
-                monitor_ccpair_permissions_taskset(tenant_id, key_bytes, r, db_session)
-
         # uncomment for debugging if needed
         # r_celery = celery_app.broker_connection().channel().client
         # length = celery_get_queue_length(DanswerCeleryQueues.VESPA_METADATA_SYNC, r_celery)
@@ -819,7 +757,7 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
 
 
 @shared_task(
-    name=DanswerCeleryTask.VESPA_METADATA_SYNC_TASK,
+    name="vespa_metadata_sync_task",
     bind=True,
     soft_time_limit=LIGHT_SOFT_TIME_LIMIT,
     time_limit=LIGHT_TIME_LIMIT,
@@ -873,9 +811,7 @@ def vespa_metadata_sync_task(
         )
     except Exception as ex:
         if isinstance(ex, RetryError):
-            task_logger.warning(
-                f"Tenacity retry failed: num_attempts={ex.last_attempt.attempt_number}"
-            )
+            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()

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

@@ -1,8 +1,6 @@
 """Factory stub for running celery worker / celery beat."""
-from celery import Celery
-
 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 = celery_app
+app = celery_app

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

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

backend/danswer/background/indexing/job_client.py

@@ -29,26 +29,18 @@ JobStatusType = (
 def _initializer(
     func: Callable, args: list | tuple, kwargs: dict[str, Any] | None = None
 ) -> Any:
-    """Initialize the child process with a fresh SQLAlchemy Engine.
+    """Ensure the parent proc's database connections are not touched
+    in the new connection pool
 
-    Based on SQLAlchemy's recommendations to handle multiprocessing:
+    Based on the recommended approach in the SQLAlchemy docs found:
     https://docs.sqlalchemy.org/en/20/core/pooling.html#using-connection-pools-with-multiprocessing-or-os-fork
     """
     if kwargs is None:
         kwargs = {}
 
     logger.info("Initializing spawned worker child process.")
-
-    # Reset the engine in the child process
-    SqlEngine.reset_engine()
-
-    # Optionally set a custom app name for database logging purposes
     SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME)
-
-    # Initialize a new engine with desired parameters
     SqlEngine.init_engine(pool_size=4, max_overflow=12, pool_recycle=60)
-
-    # Proceed with executing the target function
     return func(*args, **kwargs)
 
 

backend/danswer/background/indexing/run_indexing.py

@@ -1,5 +1,7 @@
 import time
 import traceback
+from abc import ABC
+from abc import abstractmethod
 from datetime import datetime
 from datetime import timedelta
 from datetime import timezone
@@ -19,7 +21,6 @@ from danswer.db.connector_credential_pair import get_last_successful_attempt_tim
 from danswer.db.connector_credential_pair import update_connector_credential_pair
 from danswer.db.engine import get_session_with_tenant
 from danswer.db.enums import ConnectorCredentialPairStatus
-from danswer.db.index_attempt import mark_attempt_canceled
 from danswer.db.index_attempt import mark_attempt_failed
 from danswer.db.index_attempt import mark_attempt_partially_succeeded
 from danswer.db.index_attempt import mark_attempt_succeeded
@@ -30,10 +31,10 @@ from danswer.db.models import IndexingStatus
 from danswer.db.models import IndexModelStatus
 from danswer.document_index.factory import get_default_document_index
 from danswer.indexing.embedder import DefaultIndexingEmbedder
-from danswer.indexing.indexing_heartbeat import IndexingHeartbeatInterface
+from danswer.indexing.indexing_heartbeat import IndexingHeartbeat
 from danswer.indexing.indexing_pipeline import build_indexing_pipeline
+from danswer.utils.logger import IndexAttemptSingleton
 from danswer.utils.logger import setup_logger
-from danswer.utils.logger import TaskAttemptSingleton
 from danswer.utils.variable_functionality import global_version
 
 logger = setup_logger()
@@ -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,
@@ -88,15 +102,11 @@ def _get_connector_runner(
     )
 
 
-class ConnectorStopSignal(Exception):
-    """A custom exception used to signal a stop in processing."""
-
-
 def _run_indexing(
     db_session: Session,
     index_attempt: IndexAttempt,
     tenant_id: str | None,
-    callback: IndexingHeartbeatInterface | None = None,
+    callback: RunIndexingCallbackInterface | None = None,
 ) -> None:
     """
     1. Get documents which are either new or updated from specified application
@@ -128,7 +138,13 @@ def _run_indexing(
 
     embedding_model = DefaultIndexingEmbedder.from_db_search_settings(
         search_settings=search_settings,
-        callback=callback,
+        heartbeat=IndexingHeartbeat(
+            index_attempt_id=index_attempt.id,
+            db_session=db_session,
+            # let the world know we're still making progress after
+            # every 10 batches
+            freq=10,
+        ),
     )
 
     indexing_pipeline = build_indexing_pipeline(
@@ -141,7 +157,6 @@ def _run_indexing(
         ),
         db_session=db_session,
         tenant_id=tenant_id,
-        callback=callback,
     )
 
     db_cc_pair = index_attempt.connector_credential_pair
@@ -213,7 +228,7 @@ def _run_indexing(
                 # contents still need to be initially pulled.
                 if callback:
                     if callback.should_stop():
-                        raise ConnectorStopSignal("Connector stop signal detected")
+                        raise RuntimeError("Connector stop signal detected")
 
                 # TODO: should we move this into the above callback instead?
                 db_session.refresh(db_cc_pair)
@@ -274,7 +289,7 @@ def _run_indexing(
                 db_session.commit()
 
                 if callback:
-                    callback.progress("_run_indexing", len(doc_batch))
+                    callback.progress(len(doc_batch))
 
                 # This new value is updated every batch, so UI can refresh per batch update
                 update_docs_indexed(
@@ -307,16 +322,26 @@ def _run_indexing(
                 )
         except Exception as e:
             logger.exception(
-                f"Connector run exceptioned after elapsed time: {time.time() - start_time} seconds"
+                f"Connector run ran into exception after elapsed time: {time.time() - start_time} seconds"
             )
-
-            if isinstance(e, ConnectorStopSignal):
-                mark_attempt_canceled(
+            # Only mark the attempt as a complete failure if this is the first indexing window.
+            # Otherwise, some progress was made - the next run will not start from the beginning.
+            # In this case, it is not accurate to mark it as a failure. When the next run begins,
+            # if that fails immediately, it will be marked as a failure.
+            #
+            # NOTE: if the connector is manually disabled, we should mark it as a failure regardless
+            # to give better clarity in the UI, as the next run will never happen.
+            if (
+                ind == 0
+                or not db_cc_pair.status.is_active()
+                or index_attempt.status != IndexingStatus.IN_PROGRESS
+            ):
+                mark_attempt_failed(
                     index_attempt.id,
                     db_session,
-                    reason=str(e),
+                    failure_reason=str(e),
+                    full_exception_trace=traceback.format_exc(),
                 )
-
                 if is_primary:
                     update_connector_credential_pair(
                         db_session=db_session,
@@ -328,37 +353,6 @@ def _run_indexing(
                 if INDEXING_TRACER_INTERVAL > 0:
                     tracer.stop()
                 raise e
-            else:
-                # Only mark the attempt as a complete failure if this is the first indexing window.
-                # Otherwise, some progress was made - the next run will not start from the beginning.
-                # In this case, it is not accurate to mark it as a failure. When the next run begins,
-                # if that fails immediately, it will be marked as a failure.
-                #
-                # NOTE: if the connector is manually disabled, we should mark it as a failure regardless
-                # to give better clarity in the UI, as the next run will never happen.
-                if (
-                    ind == 0
-                    or not db_cc_pair.status.is_active()
-                    or index_attempt.status != IndexingStatus.IN_PROGRESS
-                ):
-                    mark_attempt_failed(
-                        index_attempt.id,
-                        db_session,
-                        failure_reason=str(e),
-                        full_exception_trace=traceback.format_exc(),
-                    )
-
-                    if is_primary:
-                        update_connector_credential_pair(
-                            db_session=db_session,
-                            connector_id=db_connector.id,
-                            credential_id=db_credential.id,
-                            net_docs=net_doc_change,
-                        )
-
-                    if INDEXING_TRACER_INTERVAL > 0:
-                        tracer.stop()
-                    raise e
 
             # break => similar to success case. As mentioned above, if the next run fails for the same
             # reason it will then be marked as a failure
@@ -425,7 +419,7 @@ def run_indexing_entrypoint(
     tenant_id: str | None,
     connector_credential_pair_id: int,
     is_ee: bool = False,
-    callback: IndexingHeartbeatInterface | None = None,
+    callback: RunIndexingCallbackInterface | None = None,
 ) -> None:
     try:
         if is_ee:
@@ -433,19 +427,17 @@ def run_indexing_entrypoint(
 
         # set the indexing attempt ID so that all log messages from this process
         # will have it added as a prefix
-        TaskAttemptSingleton.set_cc_and_index_id(
+        IndexAttemptSingleton.set_cc_and_index_id(
             index_attempt_id, connector_credential_pair_id
         )
         with get_session_with_tenant(tenant_id) as db_session:
             attempt = transition_attempt_to_in_progress(index_attempt_id, db_session)
 
-            tenant_str = ""
-            if tenant_id is not None:
-                tenant_str = f" for tenant {tenant_id}"
-
             logger.info(
-                f"Indexing starting{tenant_str}: "
-                f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"Indexing starting for tenant {tenant_id}: "
+                if tenant_id is not None
+                else ""
+                + f"connector='{attempt.connector_credential_pair.connector.name}' "
                 f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )
@@ -453,8 +445,10 @@ def run_indexing_entrypoint(
             _run_indexing(db_session, attempt, tenant_id, callback)
 
             logger.info(
-                f"Indexing finished{tenant_str}: "
-                f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"Indexing finished for tenant {tenant_id}: "
+                if tenant_id is not None
+                else ""
+                + f"connector='{attempt.connector_credential_pair.connector.name}' "
                 f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )