fix

.github/CODEOWNERS

@@ -1,3 +1 @@
 * @onyx-dot-app/onyx-core-team
-# Helm charts Owners
-/helm/ @justin-tahara

.github/workflows/helm-chart-releases.yml

@@ -1,49 +0,0 @@
-name: Release Onyx Helm Charts
-
-on:
-  push:
-    branches:
-      - main
-
-permissions: write-all
-
-jobs:
-  release:
-    permissions:
-      contents: write
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-
-      - name: Install Helm CLI
-        uses: azure/setup-helm@v4
-        with:
-          version: v3.12.1
-
-      - name: Add required Helm repositories
-        run: |
-          helm repo add bitnami https://charts.bitnami.com/bitnami
-          helm repo add onyx-vespa https://onyx-dot-app.github.io/vespa-helm-charts
-          helm repo update
-
-      - name: Build chart dependencies
-        run: |
-          set -euo pipefail
-          for chart_dir in deployment/helm/charts/*; do
-            if [ -f "$chart_dir/Chart.yaml" ]; then
-              echo "Building dependencies for $chart_dir"
-              helm dependency build "$chart_dir"
-            fi
-          done
-
-      - name: Publish Helm charts to gh-pages
-        uses: stefanprodan/helm-gh-pages@v1.7.0
-        with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          charts_dir: deployment/helm/charts
-          branch: gh-pages
-          commit_username: ${{ github.actor }}
-          commit_email: ${{ github.actor }}@users.noreply.github.com

.github/workflows/pr-external-dependency-unit-tests.yml

@@ -13,14 +13,6 @@ env:
   # MinIO
   S3_ENDPOINT_URL: "http://localhost:9004"
 
-  # Confluence
-  CONFLUENCE_TEST_SPACE_URL: ${{ secrets.CONFLUENCE_TEST_SPACE_URL }}
-  CONFLUENCE_TEST_SPACE: ${{ secrets.CONFLUENCE_TEST_SPACE }}
-  CONFLUENCE_TEST_PAGE_ID: ${{ secrets.CONFLUENCE_TEST_PAGE_ID }}
-  CONFLUENCE_IS_CLOUD: ${{ secrets.CONFLUENCE_IS_CLOUD }}
-  CONFLUENCE_USER_NAME: ${{ secrets.CONFLUENCE_USER_NAME }}
-  CONFLUENCE_ACCESS_TOKEN: ${{ secrets.CONFLUENCE_ACCESS_TOKEN }}
-
 jobs:
   discover-test-dirs:
     runs-on: ubuntu-latest

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

@@ -55,25 +55,7 @@ jobs:
 
     - name: Run chart-testing (install)
       if: steps.list-changed.outputs.changed == 'true'
-      run: ct install --all \
-        --helm-extra-set-args="\
-          --set=nginx.enabled=false \
-          --set=postgresql.enabled=false \
-          --set=redis.enabled=false \
-          --set=minio.enabled=false \
-          --set=vespa.enabled=false \
-          --set=slackbot.enabled=false \
-          --set=api.replicaCount=0 \
-          --set=inferenceCapability.replicaCount=0 \
-          --set=indexCapability.replicaCount=0 \
-          --set=celery_beat.replicaCount=0 \
-          --set=celery_worker_heavy.replicaCount=0 \
-          --set=celery_worker_docprocessing.replicaCount=0 \
-          --set=celery_worker_light.replicaCount=0 \
-          --set=celery_worker_monitoring.replicaCount=0 \
-          --set=celery_worker_primary.replicaCount=0 \
-          --set=celery_worker_user_files_indexing.replicaCount=0" \
-        --debug --config ct.yaml
+      run: ct install --all --helm-extra-set-args="--set=nginx.enabled=false" --debug --config ct.yaml
       # the following would install only changed charts, but we only have one chart so 
       # don't worry about that for now
       # run: ct install --target-branch ${{ github.event.repository.default_branch }}

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

@@ -19,10 +19,6 @@ env:
   JIRA_BASE_URL: ${{ secrets.JIRA_BASE_URL }}
   JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}
   JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}
-  PERM_SYNC_SHAREPOINT_CLIENT_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_CLIENT_ID }}
-  PERM_SYNC_SHAREPOINT_PRIVATE_KEY: ${{ secrets.PERM_SYNC_SHAREPOINT_PRIVATE_KEY }}
-  PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD: ${{ secrets.PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD }}
-  PERM_SYNC_SHAREPOINT_DIRECTORY_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_DIRECTORY_ID }}
   PLATFORM_PAIR: linux-amd64
 
 jobs:
@@ -70,9 +66,7 @@ jobs:
             -i /local/openapi.json \
             -g python \
             -o /local/onyx_openapi_client \
-            --package-name onyx_openapi_client \
-            --skip-validate-spec \
-            --openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
+            --package-name onyx_openapi_client
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -254,8 +248,6 @@ jobs:
             -p mock-it-services-stack up -d
 
       # NOTE: Use pre-ping/null to reduce flakiness due to dropped connections
-      # NOTE: `-e ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=true` should be added once
-      # enterprise tests are fixed 
       - name: Run Standard Integration Tests
         run: |
           echo "Running integration tests..."
@@ -280,10 +272,6 @@ jobs:
             -e JIRA_BASE_URL=${JIRA_BASE_URL} \
             -e JIRA_USER_EMAIL=${JIRA_USER_EMAIL} \
             -e JIRA_API_TOKEN=${JIRA_API_TOKEN} \
-            -e PERM_SYNC_SHAREPOINT_CLIENT_ID=${PERM_SYNC_SHAREPOINT_CLIENT_ID} \
-            -e PERM_SYNC_SHAREPOINT_PRIVATE_KEY="${PERM_SYNC_SHAREPOINT_PRIVATE_KEY}" \
-            -e PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD=${PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD} \
-            -e PERM_SYNC_SHAREPOINT_DIRECTORY_ID=${PERM_SYNC_SHAREPOINT_DIRECTORY_ID} \
             -e TEST_WEB_HOSTNAME=test-runner \
             -e MOCK_CONNECTOR_SERVER_HOST=mock_connector_server \
             -e MOCK_CONNECTOR_SERVER_PORT=8001 \

.github/workflows/pr-labeler.yml

@@ -1,38 +0,0 @@
-name: PR Labeler
-
-on:
-  pull_request_target:
-    branches:
-      - main
-    types:
-      - opened
-      - reopened
-      - synchronize
-      - edited
-
-permissions:
-  contents: read
-  pull-requests: write
-
-jobs:
-  validate_pr_title:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check PR title for Conventional Commits
-        env:
-          PR_TITLE: ${{ github.event.pull_request.title }}
-        run: |
-          echo "PR Title: $PR_TITLE"
-          if [[ ! "$PR_TITLE" =~ ^(feat|fix|docs|test|ci|refactor|perf|chore|revert|build)(\(.+\))?:\ .+ ]]; then
-            echo "::error::❌ Your PR title does not follow the Conventional Commits format.
-              This check ensures that all pull requests use clear, consistent titles that help automate changelogs and improve project history.
-
-              Please update your PR title to follow the Conventional Commits style.  
-              Here is a link to a blog explaining the reason why we've included the Conventional Commits style into our PR titles: https://xfuture-blog.com/working-with-conventional-commits
-
-              **Here are some examples of valid PR titles:**
-              - feat: add user authentication
-              - fix(login): handle null password error
-              - docs(readme): update installation instructions"
-            exit 1
-          fi

.github/workflows/pr-mit-integration-tests.yml

@@ -19,10 +19,6 @@ env:
   JIRA_BASE_URL: ${{ secrets.JIRA_BASE_URL }}
   JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}
   JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}
-  PERM_SYNC_SHAREPOINT_CLIENT_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_CLIENT_ID }}
-  PERM_SYNC_SHAREPOINT_PRIVATE_KEY: ${{ secrets.PERM_SYNC_SHAREPOINT_PRIVATE_KEY }}
-  PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD: ${{ secrets.PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD }}
-  PERM_SYNC_SHAREPOINT_DIRECTORY_ID: ${{ secrets.PERM_SYNC_SHAREPOINT_DIRECTORY_ID }}
   PLATFORM_PAIR: linux-amd64
 jobs:
   integration-tests-mit:
@@ -67,9 +63,7 @@ jobs:
             -i /local/openapi.json \
             -g python \
             -o /local/onyx_openapi_client \
-            --package-name onyx_openapi_client \
-            --skip-validate-spec \
-            --openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
+            --package-name onyx_openapi_client
             
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -213,10 +207,6 @@ jobs:
             -e JIRA_BASE_URL=${JIRA_BASE_URL} \
             -e JIRA_USER_EMAIL=${JIRA_USER_EMAIL} \
             -e JIRA_API_TOKEN=${JIRA_API_TOKEN} \
-            -e PERM_SYNC_SHAREPOINT_CLIENT_ID=${PERM_SYNC_SHAREPOINT_CLIENT_ID} \
-            -e PERM_SYNC_SHAREPOINT_PRIVATE_KEY="${PERM_SYNC_SHAREPOINT_PRIVATE_KEY}" \
-            -e PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD=${PERM_SYNC_SHAREPOINT_CERTIFICATE_PASSWORD} \
-            -e PERM_SYNC_SHAREPOINT_DIRECTORY_ID=${PERM_SYNC_SHAREPOINT_DIRECTORY_ID} \
             -e TEST_WEB_HOSTNAME=test-runner \
             -e MOCK_CONNECTOR_SERVER_HOST=mock_connector_server \
             -e MOCK_CONNECTOR_SERVER_PORT=8001 \

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

@@ -47,9 +47,7 @@ jobs:
           -i /local/openapi.json \
           -g python \
           -o /local/onyx_openapi_client \
-          --package-name onyx_openapi_client \
-          --skip-validate-spec \
-          --openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
+          --package-name onyx_openapi_client
             
     - name: Run MyPy
       run: |

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

@@ -16,8 +16,8 @@ env:
   # Confluence
   CONFLUENCE_TEST_SPACE_URL: ${{ secrets.CONFLUENCE_TEST_SPACE_URL }}
   CONFLUENCE_TEST_SPACE: ${{ secrets.CONFLUENCE_TEST_SPACE }}
-  CONFLUENCE_TEST_PAGE_ID: ${{ secrets.CONFLUENCE_TEST_PAGE_ID }}
   CONFLUENCE_IS_CLOUD: ${{ secrets.CONFLUENCE_IS_CLOUD }}
+  CONFLUENCE_TEST_PAGE_ID: ${{ secrets.CONFLUENCE_TEST_PAGE_ID }}
   CONFLUENCE_USER_NAME: ${{ secrets.CONFLUENCE_USER_NAME }}
   CONFLUENCE_ACCESS_TOKEN: ${{ secrets.CONFLUENCE_ACCESS_TOKEN }}
 
@@ -53,12 +53,6 @@ env:
   # Hubspot
   HUBSPOT_ACCESS_TOKEN: ${{ secrets.HUBSPOT_ACCESS_TOKEN }}
 
-  # IMAP
-  IMAP_HOST: ${{ secrets.IMAP_HOST }}
-  IMAP_USERNAME: ${{ secrets.IMAP_USERNAME }}
-  IMAP_PASSWORD: ${{ secrets.IMAP_PASSWORD }}
-  IMAP_MAILBOXES: ${{ secrets.IMAP_MAILBOXES }}
-
   # Airtable
   AIRTABLE_TEST_BASE_ID: ${{ secrets.AIRTABLE_TEST_BASE_ID }}
   AIRTABLE_TEST_TABLE_ID: ${{ secrets.AIRTABLE_TEST_TABLE_ID }}

.gitignore

@@ -17,24 +17,12 @@ backend/tests/regression/answer_quality/test_data.json
 backend/tests/regression/search_quality/eval-*
 backend/tests/regression/search_quality/search_eval_config.yaml
 backend/tests/regression/search_quality/*.json
-*.log
 
 # secret files
 .env
 jira_test_env
-settings.json
 
 # others
 /deployment/data/nginx/app.conf
 *.sw?
 /backend/tests/regression/answer_quality/search_test_config.yaml
-
-# Local .terraform directories
-**/.terraform/*
-
-# Local .tfstate files
-*.tfstate
-*.tfstate.*
-
-# Local .terraform.lock.hcl file
-.terraform.lock.hcl

.vscode/env_template.txt

@@ -23,9 +23,6 @@ DISABLE_LLM_DOC_RELEVANCE=False
 # Useful if you want to toggle auth on/off (google_oauth/OIDC specifically)
 OAUTH_CLIENT_ID=<REPLACE THIS>
 OAUTH_CLIENT_SECRET=<REPLACE THIS>
-OPENID_CONFIG_URL=<REPLACE THIS>
-SAML_CONF_DIR=/<ABSOLUTE PATH TO ONYX>/onyx/backend/ee/onyx/configs/saml_config
-
 # Generally not useful for dev, we don't generally want to set up an SMTP server for dev
 REQUIRE_EMAIL_VERIFICATION=False
 
@@ -48,8 +45,8 @@ PYTHONPATH=../backend
 PYTHONUNBUFFERED=1
 
 
-# Internet Search
-EXA_API_KEY=<REPLACE THIS>
+# Internet Search 
+BING_API_KEY=<REPLACE THIS>
 
 
 # Enable the full set of Danswer Enterprise Edition features
@@ -67,12 +64,3 @@ S3_ENDPOINT_URL=http://localhost:9004
 S3_FILE_STORE_BUCKET_NAME=onyx-file-store-bucket
 S3_AWS_ACCESS_KEY_ID=minioadmin
 S3_AWS_SECRET_ACCESS_KEY=minioadmin
-
-# Show extra/uncommon connectors
-SHOW_EXTRA_CONNECTORS=True
-
-# Local langsmith tracing
-LANGSMITH_TRACING="true"
-LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
-LANGSMITH_API_KEY=<REPLACE_THIS>
-LANGSMITH_PROJECT=<REPLACE_THIS>

.vscode/launch.template.jsonc

@@ -24,23 +24,21 @@
           "Celery primary",
           "Celery light",
           "Celery heavy",
-          "Celery docfetching",
-          "Celery docprocessing",
+          "Celery indexing",
+          "Celery user files indexing",
           "Celery beat",
           "Celery monitoring"
         ],
         "presentation": {
           "group": "1"
-        },
-        "stopAll": true
+        }
       },
       {
         "name": "Web / Model / API",
         "configurations": ["Web Server", "Model Server", "API Server"],
         "presentation": {
           "group": "1"
-        },
-        "stopAll": true
+        }
       },
       {
         "name": "Celery (all)",
@@ -48,15 +46,14 @@
           "Celery primary",
           "Celery light",
           "Celery heavy",
-          "Celery docfetching",
-          "Celery docprocessing",
+          "Celery indexing",
+          "Celery user files indexing",
           "Celery beat",
           "Celery monitoring"
         ],
         "presentation": {
           "group": "1"
-        },
-        "stopAll": true
+        }
       }
     ],
     "configurations": [
@@ -192,7 +189,7 @@
           "--loglevel=INFO",
           "--hostname=light@%n",
           "-Q",
-          "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,index_attempt_cleanup"
+          "vespa_metadata_sync,connector_deletion,doc_permissions_upsert"
         ],
         "presentation": {
           "group": "2"
@@ -229,66 +226,35 @@
         "consoleTitle": "Celery heavy Console"
       },
       {
-        "name": "Celery docfetching",
+        "name": "Celery indexing",
         "type": "debugpy",
         "request": "launch",
         "module": "celery",
         "cwd": "${workspaceFolder}/backend",
         "envFile": "${workspaceFolder}/.vscode/.env",
         "env": {
-            "LOG_LEVEL": "DEBUG",
-            "PYTHONUNBUFFERED": "1",
-            "PYTHONPATH": "."
+          "ENABLE_MULTIPASS_INDEXING": "false",
+          "LOG_LEVEL": "DEBUG",
+          "PYTHONUNBUFFERED": "1",
+          "PYTHONPATH": "."
         },
         "args": [
-            "-A",
-            "onyx.background.celery.versioned_apps.docfetching",
-            "worker",
-            "--pool=threads",
-            "--concurrency=1",
-            "--prefetch-multiplier=1",
-            "--loglevel=INFO",
-            "--hostname=docfetching@%n",
-            "-Q",
-            "connector_doc_fetching,user_files_indexing"
+          "-A",
+          "onyx.background.celery.versioned_apps.indexing",
+          "worker",
+          "--pool=threads",
+          "--concurrency=1",
+          "--prefetch-multiplier=1",
+          "--loglevel=INFO",
+          "--hostname=indexing@%n",
+          "-Q",
+          "connector_indexing"
         ],
         "presentation": {
-            "group": "2"
+          "group": "2"
         },
-        "consoleTitle": "Celery docfetching Console",
-        "justMyCode": false
-    },
-    {
-        "name": "Celery docprocessing",
-        "type": "debugpy",
-        "request": "launch",
-        "module": "celery",
-        "cwd": "${workspaceFolder}/backend",
-        "envFile": "${workspaceFolder}/.vscode/.env",
-        "env": {
-            "ENABLE_MULTIPASS_INDEXING": "false",
-            "LOG_LEVEL": "DEBUG",
-            "PYTHONUNBUFFERED": "1",
-            "PYTHONPATH": "."
-        },
-        "args": [
-            "-A",
-            "onyx.background.celery.versioned_apps.docprocessing",
-            "worker",
-            "--pool=threads",
-            "--concurrency=6",
-            "--prefetch-multiplier=1",
-            "--loglevel=INFO",
-            "--hostname=docprocessing@%n",
-            "-Q",
-            "docprocessing"
-        ],
-        "presentation": {
-            "group": "2"
-        },
-        "consoleTitle": "Celery docprocessing Console",
-        "justMyCode": false
-    },
+        "consoleTitle": "Celery indexing Console"
+      },
       {
         "name": "Celery monitoring",
         "type": "debugpy",
@@ -337,6 +303,35 @@
         },
         "consoleTitle": "Celery beat Console"
       },
+      {
+        "name": "Celery user files indexing",
+        "type": "debugpy",
+        "request": "launch",
+        "module": "celery",
+        "cwd": "${workspaceFolder}/backend",
+        "envFile": "${workspaceFolder}/.vscode/.env",
+        "env": {
+          "LOG_LEVEL": "DEBUG",
+          "PYTHONUNBUFFERED": "1",
+          "PYTHONPATH": "."
+        },
+        "args": [
+          "-A",
+          "onyx.background.celery.versioned_apps.indexing",
+          "worker",
+          "--pool=threads",
+          "--concurrency=1",
+          "--prefetch-multiplier=1",
+          "--loglevel=INFO",
+          "--hostname=user_files_indexing@%n",
+          "-Q",
+          "user_files_indexing"
+        ],
+        "presentation": {
+          "group": "2"
+        },
+        "consoleTitle": "Celery user files indexing Console"
+      },
       {
         "name": "Pytest",
         "consoleName": "Pytest",
@@ -431,7 +426,7 @@
       },
       "args": [
         "--filename",
-        "generated/openapi.json"
+        "generated/openapi.json",
       ]
     },
     {

CONTRIBUTING.md

@@ -59,7 +59,6 @@ Onyx being a fully functional app, relies on some external software, specificall
 - [Postgres](https://www.postgresql.org/) (Relational DB)
 - [Vespa](https://vespa.ai/) (Vector DB/Search Engine)
 - [Redis](https://redis.io/) (Cache)
-- [MinIO](https://min.io/) (File Store)
 - [Nginx](https://nginx.org/) (Not needed for development flows generally)
 
 > **Note:**
@@ -103,10 +102,10 @@ If using PowerShell, the command slightly differs:
 Install the required python dependencies:
 
 ```bash
-pip install -r backend/requirements/default.txt
-pip install -r backend/requirements/dev.txt
-pip install -r backend/requirements/ee.txt
-pip install -r backend/requirements/model_server.txt
+pip install -r onyx/backend/requirements/default.txt
+pip install -r onyx/backend/requirements/dev.txt
+pip install -r onyx/backend/requirements/ee.txt
+pip install -r onyx/backend/requirements/model_server.txt
 ```
 
 Install Playwright for Python (headless browser required by the Web Connector)
@@ -172,10 +171,10 @@ Otherwise, you can follow the instructions below to run the application for deve
 
 You will need Docker installed to run these containers.
 
-First navigate to `onyx/deployment/docker_compose`, then start up Postgres/Vespa/Redis/MinIO with:
+First navigate to `onyx/deployment/docker_compose`, then start up Postgres/Vespa/Redis with:
 
 ```bash
-docker compose -f docker-compose.dev.yml -p onyx-stack up -d index relational_db cache minio
+docker compose -f docker-compose.dev.yml -p onyx-stack up -d index relational_db cache
 ```
 
 (index refers to Vespa, relational_db refers to Postgres, and cache refers to Redis)

backend/Dockerfile

@@ -12,8 +12,7 @@ ARG ONYX_VERSION=0.0.0-dev
 # DO_NOT_TRACK is used to disable telemetry for Unstructured
 ENV ONYX_VERSION=${ONYX_VERSION} \
     DANSWER_RUNNING_IN_DOCKER="true" \
-    DO_NOT_TRACK="true" \
-    PLAYWRIGHT_BROWSERS_PATH="/app/.cache/ms-playwright"
+    DO_NOT_TRACK="true"
 
 
 RUN echo "ONYX_VERSION: ${ONYX_VERSION}"
@@ -117,14 +116,6 @@ COPY ./assets /app/assets
 
 ENV PYTHONPATH=/app
 
-# Create non-root user for security best practices
-RUN groupadd -g 1001 onyx && \
-    useradd -u 1001 -g onyx -m -s /bin/bash onyx && \
-    chown -R onyx:onyx /app && \
-    mkdir -p /var/log/onyx && \
-    chmod 755 /var/log/onyx && \
-    chown onyx:onyx /var/log/onyx
-
 # Default command which does nothing
 # This container is used by api server and background which specify their own CMD
 CMD ["tail", "-f", "/dev/null"]

backend/Dockerfile.model_server

@@ -9,20 +9,11 @@ visit https://github.com/onyx-dot-app/onyx."
 # Default ONYX_VERSION, typically overriden during builds by GitHub Actions.
 ARG ONYX_VERSION=0.0.0-dev
 ENV ONYX_VERSION=${ONYX_VERSION} \
-    DANSWER_RUNNING_IN_DOCKER="true" \
-    HF_HOME=/app/.cache/huggingface
+    DANSWER_RUNNING_IN_DOCKER="true"
+
 
 RUN echo "ONYX_VERSION: ${ONYX_VERSION}"
 
-# Create non-root user for security best practices
-RUN mkdir -p /app && \
-    groupadd -g 1001 onyx && \
-    useradd -u 1001 -g onyx -m -s /bin/bash onyx  && \
-    chown -R onyx:onyx /app && \
-    mkdir -p /var/log/onyx && \
-    chmod 755 /var/log/onyx && \
-    chown onyx:onyx /var/log/onyx
-
 COPY ./requirements/model_server.txt /tmp/requirements.txt
 RUN pip install --no-cache-dir --upgrade \
         --retries 5 \
@@ -47,11 +38,9 @@ snapshot_download('mixedbread-ai/mxbai-rerank-xsmall-v1'); \
 from sentence_transformers import SentenceTransformer; \
 SentenceTransformer(model_name_or_path='nomic-ai/nomic-embed-text-v1', trust_remote_code=True);"
 
-# In case the user has volumes mounted to /app/.cache/huggingface that they've downloaded while
-# running Onyx, move the current contents of the cache folder to a temporary location to ensure 
-# it's preserved in order to combine with the user's cache contents
-RUN mv /app/.cache/huggingface /app/.cache/temp_huggingface && \
-    chown -R onyx:onyx /app
+# In case the user has volumes mounted to /root/.cache/huggingface that they've downloaded while
+# running Onyx, don't overwrite it with the built in cache folder
+RUN mv /root/.cache/huggingface /root/.cache/temp_huggingface
 
 WORKDIR /app
 

backend/alembic/env.py

@@ -23,7 +23,7 @@ from sqlalchemy.sql.schema import SchemaItem
 from onyx.configs.constants import SSL_CERT_FILE
 from shared_configs.configs import (
     MULTI_TENANT,
-    POSTGRES_DEFAULT_SCHEMA,
+    POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE,
     TENANT_ID_PREFIX,
 )
 from onyx.db.models import Base
@@ -271,7 +271,7 @@ async def run_async_migrations() -> None:
     ) = get_schema_options()
 
     if not schemas and not MULTI_TENANT:
-        schemas = [POSTGRES_DEFAULT_SCHEMA]
+        schemas = [POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE]
 
     # without init_engine, subsequent engine calls fail hard intentionally
     SqlEngine.init_engine(pool_size=20, max_overflow=5)

backend/alembic/versions/0816326d83aa_add_federated_connector_tables.py

@@ -1,72 +0,0 @@
-"""add federated connector tables
-
-Revision ID: 0816326d83aa
-Revises: 12635f6655b7
-Create Date: 2025-06-29 14:09:45.109518
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-
-# revision identifiers, used by Alembic.
-revision = "0816326d83aa"
-down_revision = "12635f6655b7"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Create federated_connector table
-    op.create_table(
-        "federated_connector",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("source", sa.String(), nullable=False),
-        sa.Column("credentials", sa.LargeBinary(), nullable=False),
-        sa.PrimaryKeyConstraint("id"),
-    )
-
-    # Create federated_connector_oauth_token table
-    op.create_table(
-        "federated_connector_oauth_token",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("federated_connector_id", sa.Integer(), nullable=False),
-        sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=False),
-        sa.Column("token", sa.LargeBinary(), nullable=False),
-        sa.Column("expires_at", sa.DateTime(), nullable=True),
-        sa.ForeignKeyConstraint(
-            ["federated_connector_id"], ["federated_connector.id"], ondelete="CASCADE"
-        ),
-        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
-        sa.PrimaryKeyConstraint("id"),
-    )
-
-    # Create federated_connector__document_set table
-    op.create_table(
-        "federated_connector__document_set",
-        sa.Column("id", sa.Integer(), nullable=False),
-        sa.Column("federated_connector_id", sa.Integer(), nullable=False),
-        sa.Column("document_set_id", sa.Integer(), nullable=False),
-        sa.Column("entities", postgresql.JSONB(), nullable=False),
-        sa.ForeignKeyConstraint(
-            ["federated_connector_id"], ["federated_connector.id"], ondelete="CASCADE"
-        ),
-        sa.ForeignKeyConstraint(
-            ["document_set_id"], ["document_set.id"], ondelete="CASCADE"
-        ),
-        sa.PrimaryKeyConstraint("id"),
-        sa.UniqueConstraint(
-            "federated_connector_id",
-            "document_set_id",
-            name="uq_federated_connector_document_set",
-        ),
-    )
-
-
-def downgrade() -> None:
-    # Drop tables in reverse order due to foreign key dependencies
-    op.drop_table("federated_connector__document_set")
-    op.drop_table("federated_connector_oauth_token")
-    op.drop_table("federated_connector")

backend/alembic/versions/12635f6655b7_drive_canonical_ids.py

@@ -1,596 +0,0 @@
-"""drive-canonical-ids
-
-Revision ID: 12635f6655b7
-Revises: 58c50ef19f08
-Create Date: 2025-06-20 14:44:54.241159
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-from urllib.parse import urlparse, urlunparse
-from httpx import HTTPStatusError
-import httpx
-from onyx.document_index.factory import get_default_document_index
-from onyx.db.search_settings import SearchSettings
-from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client
-from onyx.document_index.vespa.shared_utils.utils import (
-    replace_invalid_doc_id_characters,
-)
-from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT
-from onyx.utils.logger import setup_logger
-import os
-
-logger = setup_logger()
-
-# revision identifiers, used by Alembic.
-revision = "12635f6655b7"
-down_revision = "58c50ef19f08"
-branch_labels = None
-depends_on = None
-
-SKIP_CANON_DRIVE_IDS = os.environ.get("SKIP_CANON_DRIVE_IDS", "true").lower() == "true"
-
-
-def active_search_settings() -> tuple[SearchSettings, SearchSettings | None]:
-    result = op.get_bind().execute(
-        sa.text(
-            """
-        SELECT * FROM search_settings WHERE status = 'PRESENT' ORDER BY id DESC LIMIT 1
-        """
-        )
-    )
-    search_settings_fetch = result.fetchall()
-    search_settings = (
-        SearchSettings(**search_settings_fetch[0]._asdict())
-        if search_settings_fetch
-        else None
-    )
-
-    result2 = op.get_bind().execute(
-        sa.text(
-            """
-        SELECT * FROM search_settings WHERE status = 'FUTURE' ORDER BY id DESC LIMIT 1
-        """
-        )
-    )
-    search_settings_future_fetch = result2.fetchall()
-    search_settings_future = (
-        SearchSettings(**search_settings_future_fetch[0]._asdict())
-        if search_settings_future_fetch
-        else None
-    )
-
-    if not isinstance(search_settings, SearchSettings):
-        raise RuntimeError(
-            "current search settings is of type " + str(type(search_settings))
-        )
-    if (
-        not isinstance(search_settings_future, SearchSettings)
-        and search_settings_future is not None
-    ):
-        raise RuntimeError(
-            "future search settings is of type " + str(type(search_settings_future))
-        )
-
-    return search_settings, search_settings_future
-
-
-def normalize_google_drive_url(url: str) -> str:
-    """Remove query parameters from Google Drive URLs to create canonical document IDs.
-    NOTE: copied from drive doc_conversion.py
-    """
-    parsed_url = urlparse(url)
-    parsed_url = parsed_url._replace(query="")
-    spl_path = parsed_url.path.split("/")
-    if spl_path and (spl_path[-1] in ["edit", "view", "preview"]):
-        spl_path.pop()
-        parsed_url = parsed_url._replace(path="/".join(spl_path))
-    # Remove query parameters and reconstruct URL
-    return urlunparse(parsed_url)
-
-
-def get_google_drive_documents_from_database() -> list[dict]:
-    """Get all Google Drive documents from the database."""
-    bind = op.get_bind()
-    result = bind.execute(
-        sa.text(
-            """
-            SELECT d.id
-            FROM document d
-            JOIN document_by_connector_credential_pair dcc ON d.id = dcc.id
-            JOIN connector_credential_pair cc ON dcc.connector_id = cc.connector_id
-                AND dcc.credential_id = cc.credential_id
-            JOIN connector c ON cc.connector_id = c.id
-            WHERE c.source = 'GOOGLE_DRIVE'
-        """
-        )
-    )
-
-    documents = []
-    for row in result:
-        documents.append({"document_id": row.id})
-
-    return documents
-
-
-def update_document_id_in_database(
-    old_doc_id: str, new_doc_id: str, index_name: str
-) -> None:
-    """Update document IDs in all relevant database tables using copy-and-swap approach."""
-    bind = op.get_bind()
-
-    # print(f"Updating database tables for document {old_doc_id} -> {new_doc_id}")
-
-    # Check if new document ID already exists
-    result = bind.execute(
-        sa.text("SELECT COUNT(*) FROM document WHERE id = :new_id"),
-        {"new_id": new_doc_id},
-    )
-    row = result.fetchone()
-    if row and row[0] > 0:
-        # print(f"Document with ID {new_doc_id} already exists, deleting old one")
-        delete_document_from_db(old_doc_id, index_name)
-        return
-
-    # Step 1: Create a new document row with the new ID (copy all fields from old row)
-    # Use a conservative approach to handle columns that might not exist in all installations
-    try:
-        bind.execute(
-            sa.text(
-                """
-                INSERT INTO document (id, from_ingestion_api, boost, hidden, semantic_id,
-                                    link, doc_updated_at, primary_owners, secondary_owners,
-                                    external_user_emails, external_user_group_ids, is_public,
-                                    chunk_count, last_modified, last_synced, kg_stage, kg_processing_time)
-                SELECT :new_id, from_ingestion_api, boost, hidden, semantic_id,
-                       link, doc_updated_at, primary_owners, secondary_owners,
-                       external_user_emails, external_user_group_ids, is_public,
-                       chunk_count, last_modified, last_synced, kg_stage, kg_processing_time
-                FROM document
-                WHERE id = :old_id
-            """
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-        # print(f"Successfully updated database tables for document {old_doc_id} -> {new_doc_id}")
-    except Exception as e:
-        # If the full INSERT fails, try a more basic version with only core columns
-        logger.warning(f"Full INSERT failed, trying basic version: {e}")
-        bind.execute(
-            sa.text(
-                """
-                INSERT INTO document (id, from_ingestion_api, boost, hidden, semantic_id,
-                                    link, doc_updated_at, primary_owners, secondary_owners)
-                SELECT :new_id, from_ingestion_api, boost, hidden, semantic_id,
-                       link, doc_updated_at, primary_owners, secondary_owners
-                FROM document
-                WHERE id = :old_id
-            """
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-
-    # Step 2: Update all foreign key references to point to the new ID
-
-    # Update document_by_connector_credential_pair table
-    bind.execute(
-        sa.text(
-            "UPDATE document_by_connector_credential_pair SET id = :new_id WHERE id = :old_id"
-        ),
-        {"new_id": new_doc_id, "old_id": old_doc_id},
-    )
-    # print(f"Successfully updated document_by_connector_credential_pair table for document {old_doc_id} -> {new_doc_id}")
-
-    # Update search_doc table (stores search results for chat replay)
-    # This is critical for agent functionality
-    bind.execute(
-        sa.text(
-            "UPDATE search_doc SET document_id = :new_id WHERE document_id = :old_id"
-        ),
-        {"new_id": new_doc_id, "old_id": old_doc_id},
-    )
-    # print(f"Successfully updated search_doc table for document {old_doc_id} -> {new_doc_id}")
-    # Update document_retrieval_feedback table (user feedback on documents)
-    bind.execute(
-        sa.text(
-            "UPDATE document_retrieval_feedback SET document_id = :new_id WHERE document_id = :old_id"
-        ),
-        {"new_id": new_doc_id, "old_id": old_doc_id},
-    )
-    # print(f"Successfully updated document_retrieval_feedback table for document {old_doc_id} -> {new_doc_id}")
-    # Update document__tag table (document-tag relationships)
-    bind.execute(
-        sa.text(
-            "UPDATE document__tag SET document_id = :new_id WHERE document_id = :old_id"
-        ),
-        {"new_id": new_doc_id, "old_id": old_doc_id},
-    )
-    # print(f"Successfully updated document__tag table for document {old_doc_id} -> {new_doc_id}")
-    # Update user_file table (user uploaded files linked to documents)
-    bind.execute(
-        sa.text(
-            "UPDATE user_file SET document_id = :new_id WHERE document_id = :old_id"
-        ),
-        {"new_id": new_doc_id, "old_id": old_doc_id},
-    )
-    # print(f"Successfully updated user_file table for document {old_doc_id} -> {new_doc_id}")
-    # Update KG and chunk_stats tables (these may not exist in all installations)
-    try:
-        # Update kg_entity table
-        bind.execute(
-            sa.text(
-                "UPDATE kg_entity SET document_id = :new_id WHERE document_id = :old_id"
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-        # print(f"Successfully updated kg_entity table for document {old_doc_id} -> {new_doc_id}")
-        # Update kg_entity_extraction_staging table
-        bind.execute(
-            sa.text(
-                "UPDATE kg_entity_extraction_staging SET document_id = :new_id WHERE document_id = :old_id"
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-        # print(f"Successfully updated kg_entity_extraction_staging table for document {old_doc_id} -> {new_doc_id}")
-        # Update kg_relationship table
-        bind.execute(
-            sa.text(
-                "UPDATE kg_relationship SET source_document = :new_id WHERE source_document = :old_id"
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-        # print(f"Successfully updated kg_relationship table for document {old_doc_id} -> {new_doc_id}")
-        # Update kg_relationship_extraction_staging table
-        bind.execute(
-            sa.text(
-                "UPDATE kg_relationship_extraction_staging SET source_document = :new_id WHERE source_document = :old_id"
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-        # print(f"Successfully updated kg_relationship_extraction_staging table for document {old_doc_id} -> {new_doc_id}")
-        # Update chunk_stats table
-        bind.execute(
-            sa.text(
-                "UPDATE chunk_stats SET document_id = :new_id WHERE document_id = :old_id"
-            ),
-            {"new_id": new_doc_id, "old_id": old_doc_id},
-        )
-        # print(f"Successfully updated chunk_stats table for document {old_doc_id} -> {new_doc_id}")
-        # Update chunk_stats ID field which includes document_id
-        bind.execute(
-            sa.text(
-                """
-                UPDATE chunk_stats
-                SET id = REPLACE(id, :old_id, :new_id)
-                WHERE id LIKE :old_id_pattern
-            """
-            ),
-            {
-                "new_id": new_doc_id,
-                "old_id": old_doc_id,
-                "old_id_pattern": f"{old_doc_id}__%",
-            },
-        )
-        # print(f"Successfully updated chunk_stats ID field for document {old_doc_id} -> {new_doc_id}")
-    except Exception as e:
-        logger.warning(f"Some KG/chunk tables may not exist or failed to update: {e}")
-
-    # Step 3: Delete the old document row (this should now be safe since all FKs point to new row)
-    bind.execute(
-        sa.text("DELETE FROM document WHERE id = :old_id"), {"old_id": old_doc_id}
-    )
-    # print(f"Successfully deleted document {old_doc_id} from database")
-
-
-def _visit_chunks(
-    *,
-    http_client: httpx.Client,
-    index_name: str,
-    selection: str,
-    continuation: str | None = None,
-) -> tuple[list[dict], str | None]:
-    """Helper that calls the /document/v1 visit API once and returns (docs, next_token)."""
-
-    # Use the same URL as the document API, but with visit-specific params
-    base_url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)
-
-    params: dict[str, str] = {
-        "selection": selection,
-        "wantedDocumentCount": "1000",
-    }
-    if continuation:
-        params["continuation"] = continuation
-
-    # print(f"Visiting chunks for selection '{selection}' with params {params}")
-    resp = http_client.get(base_url, params=params, timeout=None)
-    # print(f"Visited chunks for document {selection}")
-    resp.raise_for_status()
-
-    payload = resp.json()
-    return payload.get("documents", []), payload.get("continuation")
-
-
-def delete_document_chunks_from_vespa(index_name: str, doc_id: str) -> None:
-    """Delete all chunks for *doc_id* from Vespa using continuation-token paging (no offset)."""
-
-    total_deleted = 0
-    # Use exact match instead of contains - Document Selector Language doesn't support contains
-    selection = f'{index_name}.document_id=="{doc_id}"'
-
-    with get_vespa_http_client() as http_client:
-        continuation: str | None = None
-        while True:
-            docs, continuation = _visit_chunks(
-                http_client=http_client,
-                index_name=index_name,
-                selection=selection,
-                continuation=continuation,
-            )
-
-            if not docs:
-                break
-
-            for doc in docs:
-                vespa_full_id = doc.get("id")
-                if not vespa_full_id:
-                    continue
-
-                vespa_doc_uuid = vespa_full_id.split("::")[-1]
-                delete_url = f"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{vespa_doc_uuid}"
-
-                try:
-                    resp = http_client.delete(delete_url)
-                    resp.raise_for_status()
-                    total_deleted += 1
-                except Exception as e:
-                    print(f"Failed to delete chunk {vespa_doc_uuid}: {e}")
-
-            if not continuation:
-                break
-
-
-def update_document_id_in_vespa(
-    index_name: str, old_doc_id: str, new_doc_id: str
-) -> None:
-    """Update all chunks' document_id field from *old_doc_id* to *new_doc_id* using continuation paging."""
-
-    clean_new_doc_id = replace_invalid_doc_id_characters(new_doc_id)
-
-    # Use exact match instead of contains - Document Selector Language doesn't support contains
-    selection = f'{index_name}.document_id=="{old_doc_id}"'
-
-    with get_vespa_http_client() as http_client:
-        continuation: str | None = None
-        while True:
-            # print(f"Visiting chunks for document {old_doc_id} -> {new_doc_id}")
-            docs, continuation = _visit_chunks(
-                http_client=http_client,
-                index_name=index_name,
-                selection=selection,
-                continuation=continuation,
-            )
-
-            if not docs:
-                break
-
-            for doc in docs:
-                vespa_full_id = doc.get("id")
-                if not vespa_full_id:
-                    continue
-
-                vespa_doc_uuid = vespa_full_id.split("::")[-1]
-                vespa_url = f"{DOCUMENT_ID_ENDPOINT.format(index_name=index_name)}/{vespa_doc_uuid}"
-
-                update_request = {
-                    "fields": {"document_id": {"assign": clean_new_doc_id}}
-                }
-
-                try:
-                    resp = http_client.put(vespa_url, json=update_request)
-                    resp.raise_for_status()
-                except Exception as e:
-                    print(f"Failed to update chunk {vespa_doc_uuid}: {e}")
-                    raise
-
-            if not continuation:
-                break
-
-
-def delete_document_from_db(current_doc_id: str, index_name: str) -> None:
-    # Delete all foreign key references first, then delete the document
-    try:
-        bind = op.get_bind()
-
-        # Delete from agent-related tables first (order matters due to foreign keys)
-        # Delete from agent__sub_query__search_doc first since it references search_doc
-        bind.execute(
-            sa.text(
-                """
-                DELETE FROM agent__sub_query__search_doc
-                WHERE search_doc_id IN (
-                    SELECT id FROM search_doc WHERE document_id = :doc_id
-                )
-                """
-            ),
-            {"doc_id": current_doc_id},
-        )
-
-        # Delete from chat_message__search_doc
-        bind.execute(
-            sa.text(
-                """
-                DELETE FROM chat_message__search_doc
-                WHERE search_doc_id IN (
-                    SELECT id FROM search_doc WHERE document_id = :doc_id
-                )
-                """
-            ),
-            {"doc_id": current_doc_id},
-        )
-
-        # Now we can safely delete from search_doc
-        bind.execute(
-            sa.text("DELETE FROM search_doc WHERE document_id = :doc_id"),
-            {"doc_id": current_doc_id},
-        )
-
-        # Delete from document_by_connector_credential_pair
-        bind.execute(
-            sa.text(
-                "DELETE FROM document_by_connector_credential_pair WHERE id = :doc_id"
-            ),
-            {"doc_id": current_doc_id},
-        )
-
-        # Delete from other tables that reference this document
-        bind.execute(
-            sa.text(
-                "DELETE FROM document_retrieval_feedback WHERE document_id = :doc_id"
-            ),
-            {"doc_id": current_doc_id},
-        )
-
-        bind.execute(
-            sa.text("DELETE FROM document__tag WHERE document_id = :doc_id"),
-            {"doc_id": current_doc_id},
-        )
-
-        bind.execute(
-            sa.text("DELETE FROM user_file WHERE document_id = :doc_id"),
-            {"doc_id": current_doc_id},
-        )
-
-        # Delete from KG tables if they exist
-        try:
-            bind.execute(
-                sa.text("DELETE FROM kg_entity WHERE document_id = :doc_id"),
-                {"doc_id": current_doc_id},
-            )
-
-            bind.execute(
-                sa.text(
-                    "DELETE FROM kg_entity_extraction_staging WHERE document_id = :doc_id"
-                ),
-                {"doc_id": current_doc_id},
-            )
-
-            bind.execute(
-                sa.text("DELETE FROM kg_relationship WHERE source_document = :doc_id"),
-                {"doc_id": current_doc_id},
-            )
-
-            bind.execute(
-                sa.text(
-                    "DELETE FROM kg_relationship_extraction_staging WHERE source_document = :doc_id"
-                ),
-                {"doc_id": current_doc_id},
-            )
-
-            bind.execute(
-                sa.text("DELETE FROM chunk_stats WHERE document_id = :doc_id"),
-                {"doc_id": current_doc_id},
-            )
-
-            bind.execute(
-                sa.text("DELETE FROM chunk_stats WHERE id LIKE :doc_id_pattern"),
-                {"doc_id_pattern": f"{current_doc_id}__%"},
-            )
-
-        except Exception as e:
-            logger.warning(
-                f"Some KG/chunk tables may not exist or failed to delete from: {e}"
-            )
-
-        # Finally delete the document itself
-        bind.execute(
-            sa.text("DELETE FROM document WHERE id = :doc_id"),
-            {"doc_id": current_doc_id},
-        )
-
-        # Delete chunks from vespa
-        delete_document_chunks_from_vespa(index_name, current_doc_id)
-
-    except Exception as e:
-        print(f"Failed to delete duplicate document {current_doc_id}: {e}")
-        # Continue with other documents instead of failing the entire migration
-
-
-def upgrade() -> None:
-    if SKIP_CANON_DRIVE_IDS:
-        return
-    current_search_settings, future_search_settings = active_search_settings()
-    document_index = get_default_document_index(
-        current_search_settings,
-        future_search_settings,
-    )
-
-    # Get the index name
-    if hasattr(document_index, "index_name"):
-        index_name = document_index.index_name
-    else:
-        # Default index name if we can't get it from the document_index
-        index_name = "danswer_index"
-
-    # Get all Google Drive documents from the database (this is faster and more reliable)
-    gdrive_documents = get_google_drive_documents_from_database()
-
-    if not gdrive_documents:
-        return
-
-    # Track normalized document IDs to detect duplicates
-    all_normalized_doc_ids = set()
-    updated_count = 0
-
-    for doc_info in gdrive_documents:
-        current_doc_id = doc_info["document_id"]
-        normalized_doc_id = normalize_google_drive_url(current_doc_id)
-
-        print(f"Processing document {current_doc_id} -> {normalized_doc_id}")
-        # Check for duplicates
-        if normalized_doc_id in all_normalized_doc_ids:
-            # print(f"Deleting duplicate document {current_doc_id}")
-            delete_document_from_db(current_doc_id, index_name)
-            continue
-
-        all_normalized_doc_ids.add(normalized_doc_id)
-
-        # If the document ID already doesn't have query parameters, skip it
-        if current_doc_id == normalized_doc_id:
-            # print(f"Skipping document {current_doc_id} -> {normalized_doc_id} because it already has no query parameters")
-            continue
-
-        try:
-            # Update both database and Vespa in order
-            # Database first to ensure consistency
-            update_document_id_in_database(
-                current_doc_id, normalized_doc_id, index_name
-            )
-
-            # For Vespa, we can now use the original document IDs since we're using contains matching
-            update_document_id_in_vespa(index_name, current_doc_id, normalized_doc_id)
-            updated_count += 1
-            # print(f"Finished updating document {current_doc_id} -> {normalized_doc_id}")
-        except Exception as e:
-            print(f"Failed to update document {current_doc_id}: {e}")
-
-            if isinstance(e, HTTPStatusError):
-                print(f"HTTPStatusError: {e}")
-                print(f"Response: {e.response.text}")
-                print(f"Status: {e.response.status_code}")
-                print(f"Headers: {e.response.headers}")
-                print(f"Request: {e.request.url}")
-                print(f"Request headers: {e.request.headers}")
-            # Note: Rollback is complex with copy-and-swap approach since the old document is already deleted
-            # In case of failure, manual intervention may be required
-            # Continue with other documents instead of failing the entire migration
-            continue
-
-    logger.info(f"Migration complete. Updated {updated_count} Google Drive documents")
-
-
-def downgrade() -> None:
-    # this is a one way migration, so no downgrade.
-    # It wouldn't make sense to store the extra query parameters
-    # and duplicate documents to allow a reversal.
-    pass

backend/alembic/versions/27c6ecc08586_permission_framework.py

@@ -144,34 +144,27 @@ def upgrade() -> None:
 
 def downgrade() -> None:
     op.execute("TRUNCATE TABLE index_attempt")
-    conn = op.get_bind()
-    inspector = sa.inspect(conn)
-    existing_columns = {col["name"] for col in inspector.get_columns("index_attempt")}
-
-    if "input_type" not in existing_columns:
-        op.add_column(
-            "index_attempt",
-            sa.Column("input_type", sa.VARCHAR(), autoincrement=False, nullable=False),
-        )
-
-    if "source" not in existing_columns:
-        op.add_column(
-            "index_attempt",
-            sa.Column("source", sa.VARCHAR(), autoincrement=False, nullable=False),
-        )
-
-    if "connector_specific_config" not in existing_columns:
-        op.add_column(
-            "index_attempt",
-            sa.Column(
-                "connector_specific_config",
-                postgresql.JSONB(astext_type=sa.Text()),
-                autoincrement=False,
-                nullable=False,
-            ),
-        )
+    op.add_column(
+        "index_attempt",
+        sa.Column("input_type", sa.VARCHAR(), autoincrement=False, nullable=False),
+    )
+    op.add_column(
+        "index_attempt",
+        sa.Column("source", sa.VARCHAR(), autoincrement=False, nullable=False),
+    )
+    op.add_column(
+        "index_attempt",
+        sa.Column(
+            "connector_specific_config",
+            postgresql.JSONB(astext_type=sa.Text()),
+            autoincrement=False,
+            nullable=False,
+        ),
+    )
 
     # Check if the constraint exists before dropping
+    conn = op.get_bind()
+    inspector = sa.inspect(conn)
     constraints = inspector.get_foreign_keys("index_attempt")
 
     if any(
@@ -190,12 +183,8 @@ def downgrade() -> None:
             "fk_index_attempt_connector_id", "index_attempt", type_="foreignkey"
         )
 
-    if "credential_id" in existing_columns:
-        op.drop_column("index_attempt", "credential_id")
-
-    if "connector_id" in existing_columns:
-        op.drop_column("index_attempt", "connector_id")
-
-    op.execute("DROP TABLE IF EXISTS connector_credential_pair CASCADE")
-    op.execute("DROP TABLE IF EXISTS credential CASCADE")
-    op.execute("DROP TABLE IF EXISTS connector CASCADE")
+    op.drop_column("index_attempt", "credential_id")
+    op.drop_column("index_attempt", "connector_id")
+    op.drop_table("connector_credential_pair")
+    op.drop_table("credential")
+    op.drop_table("connector")

backend/alembic/versions/2f95e36923e6_add_indexing_coordination.py

@@ -1,115 +0,0 @@
-"""add_indexing_coordination
-
-Revision ID: 2f95e36923e6
-Revises: 0816326d83aa
-Create Date: 2025-07-10 16:17:57.762182
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "2f95e36923e6"
-down_revision = "0816326d83aa"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Add database-based coordination fields (replacing Redis fencing)
-    op.add_column(
-        "index_attempt", sa.Column("celery_task_id", sa.String(), nullable=True)
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column(
-            "cancellation_requested",
-            sa.Boolean(),
-            nullable=False,
-            server_default="false",
-        ),
-    )
-
-    # Add batch coordination fields (replacing FileStore state)
-    op.add_column(
-        "index_attempt", sa.Column("total_batches", sa.Integer(), nullable=True)
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column(
-            "completed_batches", sa.Integer(), nullable=False, server_default="0"
-        ),
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column(
-            "total_failures_batch_level",
-            sa.Integer(),
-            nullable=False,
-            server_default="0",
-        ),
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column("total_chunks", sa.Integer(), nullable=False, server_default="0"),
-    )
-
-    # Progress tracking for stall detection
-    op.add_column(
-        "index_attempt",
-        sa.Column("last_progress_time", sa.DateTime(timezone=True), nullable=True),
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column(
-            "last_batches_completed_count",
-            sa.Integer(),
-            nullable=False,
-            server_default="0",
-        ),
-    )
-
-    # Heartbeat tracking for worker liveness detection
-    op.add_column(
-        "index_attempt",
-        sa.Column(
-            "heartbeat_counter", sa.Integer(), nullable=False, server_default="0"
-        ),
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column(
-            "last_heartbeat_value", sa.Integer(), nullable=False, server_default="0"
-        ),
-    )
-    op.add_column(
-        "index_attempt",
-        sa.Column("last_heartbeat_time", sa.DateTime(timezone=True), nullable=True),
-    )
-
-    # Add index for coordination queries
-    op.create_index(
-        "ix_index_attempt_active_coordination",
-        "index_attempt",
-        ["connector_credential_pair_id", "search_settings_id", "status"],
-    )
-
-
-def downgrade() -> None:
-    # Remove the new index
-    op.drop_index("ix_index_attempt_active_coordination", table_name="index_attempt")
-
-    # Remove the new columns
-    op.drop_column("index_attempt", "last_batches_completed_count")
-    op.drop_column("index_attempt", "last_progress_time")
-    op.drop_column("index_attempt", "last_heartbeat_time")
-    op.drop_column("index_attempt", "last_heartbeat_value")
-    op.drop_column("index_attempt", "heartbeat_counter")
-    op.drop_column("index_attempt", "total_chunks")
-    op.drop_column("index_attempt", "total_failures_batch_level")
-    op.drop_column("index_attempt", "completed_batches")
-    op.drop_column("index_attempt", "total_batches")
-    op.drop_column("index_attempt", "cancellation_requested")
-    op.drop_column("index_attempt", "celery_task_id")

backend/alembic/versions/36e9220ab794_update_kg_trigger_functions.py

@@ -9,7 +9,7 @@ Create Date: 2025-06-22 17:33:25.833733
 from alembic import op
 from sqlalchemy.orm import Session
 from sqlalchemy import text
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE
 
 # revision identifiers, used by Alembic.
 revision = "36e9220ab794"
@@ -66,7 +66,7 @@ def upgrade() -> None:
 
                 -- Set name and name trigrams
                 NEW.name = name;
-                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name);
+                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.show_trgm(cleaned_name);
                 RETURN NEW;
             END;
             $$ LANGUAGE plpgsql;
@@ -111,7 +111,7 @@ def upgrade() -> None:
                 UPDATE "{tenant_id}".kg_entity
                 SET
                     name = doc_name,
-                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name)
+                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.show_trgm(cleaned_name)
                 WHERE document_id = NEW.id;
                 RETURN NEW;
             END;

backend/alembic/versions/3fc5d75723b3_add_doc_metadata_field_in_document_model.py

@@ -1,30 +0,0 @@
-"""add_doc_metadata_field_in_document_model
-
-Revision ID: 3fc5d75723b3
-Revises: 2f95e36923e6
-Create Date: 2025-07-28 18:45:37.985406
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-# revision identifiers, used by Alembic.
-revision = "3fc5d75723b3"
-down_revision = "2f95e36923e6"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.add_column(
-        "document",
-        sa.Column(
-            "doc_metadata", postgresql.JSONB(astext_type=sa.Text()), nullable=True
-        ),
-    )
-
-
-def downgrade() -> None:
-    op.drop_column("document", "doc_metadata")

backend/alembic/versions/495cb26ce93e_create_knowlege_graph_tables.py

@@ -15,7 +15,7 @@ from datetime import datetime, timedelta
 from onyx.configs.app_configs import DB_READONLY_USER
 from onyx.configs.app_configs import DB_READONLY_PASSWORD
 from shared_configs.configs import MULTI_TENANT
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE
 
 
 # revision identifiers, used by Alembic.
@@ -80,7 +80,6 @@ def upgrade() -> None:
         )
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_config CASCADE")
     op.create_table(
         "kg_config",
         sa.Column("id", sa.Integer(), primary_key=True, nullable=False, index=True),
@@ -124,7 +123,6 @@ def upgrade() -> None:
         ],
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_entity_type CASCADE")
     op.create_table(
         "kg_entity_type",
         sa.Column("id_name", sa.String(), primary_key=True, nullable=False, index=True),
@@ -158,7 +156,6 @@ def upgrade() -> None:
         ),
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_relationship_type CASCADE")
     # Create KGRelationshipType table
     op.create_table(
         "kg_relationship_type",
@@ -197,7 +194,6 @@ def upgrade() -> None:
         ),
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_relationship_type_extraction_staging CASCADE")
     # Create KGRelationshipTypeExtractionStaging table
     op.create_table(
         "kg_relationship_type_extraction_staging",
@@ -231,8 +227,6 @@ def upgrade() -> None:
         ),
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_entity CASCADE")
-
     # Create KGEntity table
     op.create_table(
         "kg_entity",
@@ -287,7 +281,6 @@ def upgrade() -> None:
         "ix_entity_name_search", "kg_entity", ["name", "entity_type_id_name"]
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_entity_extraction_staging CASCADE")
     # Create KGEntityExtractionStaging table
     op.create_table(
         "kg_entity_extraction_staging",
@@ -337,7 +330,6 @@ def upgrade() -> None:
         ["name", "entity_type_id_name"],
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_relationship CASCADE")
     # Create KGRelationship table
     op.create_table(
         "kg_relationship",
@@ -379,7 +371,6 @@ def upgrade() -> None:
         "ix_kg_relationship_nodes", "kg_relationship", ["source_node", "target_node"]
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_relationship_extraction_staging CASCADE")
     # Create KGRelationshipExtractionStaging table
     op.create_table(
         "kg_relationship_extraction_staging",
@@ -423,7 +414,6 @@ def upgrade() -> None:
         ["source_node", "target_node"],
     )
 
-    op.execute("DROP TABLE IF EXISTS kg_term CASCADE")
     # Create KGTerm table
     op.create_table(
         "kg_term",
@@ -478,7 +468,7 @@ def upgrade() -> None:
     # Create GIN index for clustering and normalization
     op.execute(
         "CREATE INDEX IF NOT EXISTS idx_kg_entity_clustering_trigrams "
-        f"ON kg_entity USING GIN (name {POSTGRES_DEFAULT_SCHEMA}.gin_trgm_ops)"
+        f"ON kg_entity USING GIN (name {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.gin_trgm_ops)"
     )
     op.execute(
         "CREATE INDEX IF NOT EXISTS idx_kg_entity_normalization_trigrams "
@@ -518,7 +508,7 @@ def upgrade() -> None:
 
                 -- Set name and name trigrams
                 NEW.name = name;
-                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name);
+                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.show_trgm(cleaned_name);
                 RETURN NEW;
             END;
             $$ LANGUAGE plpgsql;
@@ -563,7 +553,7 @@ def upgrade() -> None:
                 UPDATE kg_entity
                 SET
                     name = doc_name,
-                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name)
+                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.show_trgm(cleaned_name)
                 WHERE document_id = NEW.id;
                 RETURN NEW;
             END;

backend/alembic/versions/5ae8240accb3_add_research_agent_database_tables_and_.py

@@ -1,115 +0,0 @@
-"""add research agent database tables and chat message research fields
-
-Revision ID: 5ae8240accb3
-Revises: b558f51620b4
-Create Date: 2025-08-06 14:29:24.691388
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-
-# revision identifiers, used by Alembic.
-revision = "5ae8240accb3"
-down_revision = "b558f51620b4"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Add research_type and research_plan columns to chat_message table
-    op.add_column(
-        "chat_message",
-        sa.Column("research_type", sa.String(), nullable=True),
-    )
-    op.add_column(
-        "chat_message",
-        sa.Column("research_plan", postgresql.JSONB(), nullable=True),
-    )
-
-    # Create research_agent_iteration table
-    op.create_table(
-        "research_agent_iteration",
-        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
-        sa.Column(
-            "primary_question_id",
-            sa.Integer(),
-            sa.ForeignKey("chat_message.id", ondelete="CASCADE"),
-            nullable=False,
-        ),
-        sa.Column("iteration_nr", sa.Integer(), nullable=False),
-        sa.Column(
-            "created_at",
-            sa.DateTime(timezone=True),
-            server_default=sa.func.now(),
-            nullable=False,
-        ),
-        sa.Column("purpose", sa.String(), nullable=True),
-        sa.Column("reasoning", sa.String(), nullable=True),
-        sa.PrimaryKeyConstraint("id"),
-        sa.UniqueConstraint(
-            "primary_question_id",
-            "iteration_nr",
-            name="_research_agent_iteration_unique_constraint",
-        ),
-    )
-
-    # Create research_agent_iteration_sub_step table
-    op.create_table(
-        "research_agent_iteration_sub_step",
-        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
-        sa.Column(
-            "primary_question_id",
-            sa.Integer(),
-            sa.ForeignKey("chat_message.id", ondelete="CASCADE"),
-            nullable=False,
-        ),
-        sa.Column(
-            "parent_question_id",
-            sa.Integer(),
-            sa.ForeignKey("research_agent_iteration_sub_step.id", ondelete="CASCADE"),
-            nullable=True,
-        ),
-        sa.Column("iteration_nr", sa.Integer(), nullable=False),
-        sa.Column("iteration_sub_step_nr", sa.Integer(), nullable=False),
-        sa.Column(
-            "created_at",
-            sa.DateTime(timezone=True),
-            server_default=sa.func.now(),
-            nullable=False,
-        ),
-        sa.Column("sub_step_instructions", sa.String(), nullable=True),
-        sa.Column(
-            "sub_step_tool_id",
-            sa.Integer(),
-            sa.ForeignKey("tool.id"),
-            nullable=True,
-        ),
-        sa.Column("reasoning", sa.String(), nullable=True),
-        sa.Column("sub_answer", sa.String(), nullable=True),
-        sa.Column("cited_doc_results", postgresql.JSONB(), nullable=True),
-        sa.Column("claims", postgresql.JSONB(), nullable=True),
-        sa.Column("generated_images", postgresql.JSONB(), nullable=True),
-        sa.Column("additional_data", postgresql.JSONB(), nullable=True),
-        sa.PrimaryKeyConstraint("id"),
-        sa.ForeignKeyConstraint(
-            ["primary_question_id", "iteration_nr"],
-            [
-                "research_agent_iteration.primary_question_id",
-                "research_agent_iteration.iteration_nr",
-            ],
-            ondelete="CASCADE",
-        ),
-    )
-
-
-def downgrade() -> None:
-    # Drop tables in reverse order
-    op.drop_table("research_agent_iteration_sub_step")
-    op.drop_table("research_agent_iteration")
-
-    # Remove columns from chat_message table
-    op.drop_column("chat_message", "research_plan")
-    op.drop_column("chat_message", "research_type")

backend/alembic/versions/62c3a055a141_add_file_names_to_file_connector_config.py

@@ -1,132 +0,0 @@
-"""add file names to file connector config
-
-Revision ID: 62c3a055a141
-Revises: 3fc5d75723b3
-Create Date: 2025-07-30 17:01:24.417551
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-import json
-import os
-import logging
-
-
-# revision identifiers, used by Alembic.
-revision = "62c3a055a141"
-down_revision = "3fc5d75723b3"
-branch_labels = None
-depends_on = None
-
-SKIP_FILE_NAME_MIGRATION = (
-    os.environ.get("SKIP_FILE_NAME_MIGRATION", "true").lower() == "true"
-)
-
-logger = logging.getLogger("alembic.runtime.migration")
-
-
-def upgrade() -> None:
-    if SKIP_FILE_NAME_MIGRATION:
-        logger.info(
-            "Skipping file name migration. Hint: set SKIP_FILE_NAME_MIGRATION=false to run this migration"
-        )
-        return
-    logger.info("Running file name migration")
-    # Get connection
-    conn = op.get_bind()
-
-    # Get all FILE connectors with their configs
-    file_connectors = conn.execute(
-        sa.text(
-            """
-            SELECT id, connector_specific_config
-            FROM connector
-            WHERE source = 'FILE'
-        """
-        )
-    ).fetchall()
-
-    for connector_id, config in file_connectors:
-        # Parse config if it's a string
-        if isinstance(config, str):
-            config = json.loads(config)
-
-        # Get file_locations list
-        file_locations = config.get("file_locations", [])
-
-        # Get display names for each file_id
-        file_names = []
-        for file_id in file_locations:
-            result = conn.execute(
-                sa.text(
-                    """
-                    SELECT display_name
-                    FROM file_record
-                    WHERE file_id = :file_id
-                """
-                ),
-                {"file_id": file_id},
-            ).fetchone()
-
-            if result:
-                file_names.append(result[0])
-            else:
-                file_names.append(file_id)  # Should not happen
-
-        # Add file_names to config
-        new_config = dict(config)
-        new_config["file_names"] = file_names
-
-        # Update the connector
-        conn.execute(
-            sa.text(
-                """
-                UPDATE connector
-                SET connector_specific_config = :new_config
-                WHERE id = :connector_id
-            """
-            ),
-            {"connector_id": connector_id, "new_config": json.dumps(new_config)},
-        )
-
-
-def downgrade() -> None:
-    # Get connection
-    conn = op.get_bind()
-
-    # Remove file_names from all FILE connectors
-    file_connectors = conn.execute(
-        sa.text(
-            """
-            SELECT id, connector_specific_config
-            FROM connector
-            WHERE source = 'FILE'
-        """
-        )
-    ).fetchall()
-
-    for connector_id, config in file_connectors:
-        # Parse config if it's a string
-        if isinstance(config, str):
-            config = json.loads(config)
-
-        # Remove file_names if it exists
-        if "file_names" in config:
-            new_config = dict(config)
-            del new_config["file_names"]
-
-            # Update the connector
-            conn.execute(
-                sa.text(
-                    """
-                    UPDATE connector
-                    SET connector_specific_config = :new_config
-                    WHERE id = :connector_id
-                """
-                ),
-                {
-                    "connector_id": connector_id,
-                    "new_config": json.dumps(new_config),
-                },
-            )

backend/alembic/versions/7ed603b64d5a_add_mcp_server_and_connection_config_.py

@@ -1,249 +0,0 @@
-"""add_mcp_server_and_connection_config_models
-
-Revision ID: 7ed603b64d5a
-Revises: b329d00a9ea6
-Create Date: 2025-07-28 17:35:59.900680
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-from onyx.db.enums import MCPAuthenticationType
-
-# revision identifiers, used by Alembic.
-revision = "7ed603b64d5a"
-down_revision = "b329d00a9ea6"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    """Create tables and columns for MCP Server support"""
-
-    # 1. MCP Server main table (no FK constraints yet to avoid circular refs)
-    op.create_table(
-        "mcp_server",
-        sa.Column("id", sa.Integer(), primary_key=True),
-        sa.Column("owner", sa.String(), nullable=False),
-        sa.Column("name", sa.String(), nullable=False),
-        sa.Column("description", sa.String(), nullable=True),
-        sa.Column("server_url", sa.String(), nullable=False),
-        sa.Column(
-            "auth_type",
-            sa.Enum(
-                MCPAuthenticationType,
-                name="mcp_authentication_type",
-                native_enum=False,
-            ),
-            nullable=False,
-        ),
-        sa.Column("admin_connection_config_id", sa.Integer(), nullable=True),
-        sa.Column(
-            "created_at",
-            sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
-            nullable=False,
-        ),
-        sa.Column(
-            "updated_at",
-            sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
-            nullable=False,
-        ),
-    )
-
-    # 2. MCP Connection Config table (can reference mcp_server now that it exists)
-    op.create_table(
-        "mcp_connection_config",
-        sa.Column("id", sa.Integer(), primary_key=True),
-        sa.Column("mcp_server_id", sa.Integer(), nullable=True),
-        sa.Column("user_email", sa.String(), nullable=False, default=""),
-        sa.Column("config", sa.LargeBinary(), nullable=False),
-        sa.Column(
-            "created_at",
-            sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
-            nullable=False,
-        ),
-        sa.Column(
-            "updated_at",
-            sa.DateTime(timezone=True),
-            server_default=sa.text("now()"),  # type: ignore
-            nullable=False,
-        ),
-        sa.ForeignKeyConstraint(
-            ["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
-        ),
-    )
-
-    # Helpful indexes
-    op.create_index(
-        "ix_mcp_connection_config_server_user",
-        "mcp_connection_config",
-        ["mcp_server_id", "user_email"],
-    )
-    op.create_index(
-        "ix_mcp_connection_config_user_email",
-        "mcp_connection_config",
-        ["user_email"],
-    )
-
-    # 3. Add the back-references from mcp_server to connection configs
-    op.create_foreign_key(
-        "mcp_server_admin_config_fk",
-        "mcp_server",
-        "mcp_connection_config",
-        ["admin_connection_config_id"],
-        ["id"],
-        ondelete="SET NULL",
-    )
-
-    # 4. Association / access-control tables
-    op.create_table(
-        "mcp_server__user",
-        sa.Column("mcp_server_id", sa.Integer(), primary_key=True),
-        sa.Column("user_id", sa.UUID(), primary_key=True),
-        sa.ForeignKeyConstraint(
-            ["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
-        ),
-        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
-    )
-
-    op.create_table(
-        "mcp_server__user_group",
-        sa.Column("mcp_server_id", sa.Integer(), primary_key=True),
-        sa.Column("user_group_id", sa.Integer(), primary_key=True),
-        sa.ForeignKeyConstraint(
-            ["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE"
-        ),
-        sa.ForeignKeyConstraint(["user_group_id"], ["user_group.id"]),
-    )
-
-    # 5. Update existing `tool` table – allow tools to belong to an MCP server
-    op.add_column(
-        "tool",
-        sa.Column("mcp_server_id", sa.Integer(), nullable=True),
-    )
-    # Add column for MCP tool input schema
-    op.add_column(
-        "tool",
-        sa.Column("mcp_input_schema", postgresql.JSONB(), nullable=True),
-    )
-    op.create_foreign_key(
-        "tool_mcp_server_fk",
-        "tool",
-        "mcp_server",
-        ["mcp_server_id"],
-        ["id"],
-        ondelete="CASCADE",
-    )
-
-    # 6. Update persona__tool foreign keys to cascade delete
-    # This ensures that when a tool is deleted (including via MCP server deletion),
-    # the corresponding persona__tool rows are also deleted
-    op.drop_constraint(
-        "persona__tool_tool_id_fkey", "persona__tool", type_="foreignkey"
-    )
-    op.drop_constraint(
-        "persona__tool_persona_id_fkey", "persona__tool", type_="foreignkey"
-    )
-
-    op.create_foreign_key(
-        "persona__tool_persona_id_fkey",
-        "persona__tool",
-        "persona",
-        ["persona_id"],
-        ["id"],
-        ondelete="CASCADE",
-    )
-    op.create_foreign_key(
-        "persona__tool_tool_id_fkey",
-        "persona__tool",
-        "tool",
-        ["tool_id"],
-        ["id"],
-        ondelete="CASCADE",
-    )
-
-    # 7. Update research_agent_iteration_sub_step foreign key to SET NULL on delete
-    # This ensures that when a tool is deleted, the sub_step_tool_id is set to NULL
-    # instead of causing a foreign key constraint violation
-    op.drop_constraint(
-        "research_agent_iteration_sub_step_sub_step_tool_id_fkey",
-        "research_agent_iteration_sub_step",
-        type_="foreignkey",
-    )
-    op.create_foreign_key(
-        "research_agent_iteration_sub_step_sub_step_tool_id_fkey",
-        "research_agent_iteration_sub_step",
-        "tool",
-        ["sub_step_tool_id"],
-        ["id"],
-        ondelete="SET NULL",
-    )
-
-
-def downgrade() -> None:
-    """Drop all MCP-related tables / columns"""
-
-    # # # 1. Drop FK & columns from tool
-    # op.drop_constraint("tool_mcp_server_fk", "tool", type_="foreignkey")
-    op.execute("DELETE FROM tool WHERE mcp_server_id IS NOT NULL")
-
-    op.drop_constraint(
-        "research_agent_iteration_sub_step_sub_step_tool_id_fkey",
-        "research_agent_iteration_sub_step",
-        type_="foreignkey",
-    )
-    op.create_foreign_key(
-        "research_agent_iteration_sub_step_sub_step_tool_id_fkey",
-        "research_agent_iteration_sub_step",
-        "tool",
-        ["sub_step_tool_id"],
-        ["id"],
-    )
-
-    # Restore original persona__tool foreign keys (without CASCADE)
-    op.drop_constraint(
-        "persona__tool_persona_id_fkey", "persona__tool", type_="foreignkey"
-    )
-    op.drop_constraint(
-        "persona__tool_tool_id_fkey", "persona__tool", type_="foreignkey"
-    )
-
-    op.create_foreign_key(
-        "persona__tool_persona_id_fkey",
-        "persona__tool",
-        "persona",
-        ["persona_id"],
-        ["id"],
-    )
-    op.create_foreign_key(
-        "persona__tool_tool_id_fkey",
-        "persona__tool",
-        "tool",
-        ["tool_id"],
-        ["id"],
-    )
-    op.drop_column("tool", "mcp_input_schema")
-    op.drop_column("tool", "mcp_server_id")
-
-    # 2. Drop association tables
-    op.drop_table("mcp_server__user_group")
-    op.drop_table("mcp_server__user")
-
-    # 3. Drop FK from mcp_server to connection configs
-    op.drop_constraint("mcp_server_admin_config_fk", "mcp_server", type_="foreignkey")
-
-    # 4. Drop connection config indexes & table
-    op.drop_index(
-        "ix_mcp_connection_config_user_email", table_name="mcp_connection_config"
-    )
-    op.drop_index(
-        "ix_mcp_connection_config_server_user", table_name="mcp_connection_config"
-    )
-    op.drop_table("mcp_connection_config")
-
-    # 5. Finally drop mcp_server table
-    op.drop_table("mcp_server")

backend/alembic/versions/8818cf73fa1a_drop_include_citations.py

@@ -1,38 +0,0 @@
-"""drop include citations
-
-Revision ID: 8818cf73fa1a
-Revises: 7ed603b64d5a
-Create Date: 2025-09-02 19:43:50.060680
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-
-# revision identifiers, used by Alembic.
-revision = "8818cf73fa1a"
-down_revision = "7ed603b64d5a"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.drop_column("prompt", "include_citations")
-
-
-def downgrade() -> None:
-    op.add_column(
-        "prompt",
-        sa.Column(
-            "include_citations",
-            sa.BOOLEAN(),
-            autoincrement=False,
-            nullable=True,
-        ),
-    )
-    # Set include_citations based on prompt name: FALSE for ImageGeneration, TRUE for others
-    op.execute(
-        sa.text(
-            "UPDATE prompt SET include_citations = CASE WHEN name = 'ImageGeneration' THEN FALSE ELSE TRUE END"
-        )
-    )

backend/alembic/versions/90e3b9af7da4_tag_fix.py

@@ -1,341 +0,0 @@
-"""tag-fix
-
-Revision ID: 90e3b9af7da4
-Revises: 62c3a055a141
-Create Date: 2025-08-01 20:58:14.607624
-
-"""
-
-import json
-import logging
-import os
-
-from typing import cast
-from typing import Generator
-
-from alembic import op
-import sqlalchemy as sa
-
-from onyx.document_index.factory import get_default_document_index
-from onyx.document_index.vespa_constants import DOCUMENT_ID_ENDPOINT
-from onyx.db.search_settings import SearchSettings
-from onyx.configs.app_configs import AUTH_TYPE
-from onyx.configs.constants import AuthType
-from onyx.document_index.vespa.shared_utils.utils import get_vespa_http_client
-
-logger = logging.getLogger("alembic.runtime.migration")
-
-
-# revision identifiers, used by Alembic.
-revision = "90e3b9af7da4"
-down_revision = "62c3a055a141"
-branch_labels = None
-depends_on = None
-
-SKIP_TAG_FIX = os.environ.get("SKIP_TAG_FIX", "true").lower() == "true"
-
-# override for cloud
-if AUTH_TYPE == AuthType.CLOUD:
-    SKIP_TAG_FIX = True
-
-
-def set_is_list_for_known_tags() -> None:
-    """
-    Sets is_list to true for all tags that are known to be lists.
-    """
-    LIST_METADATA: list[tuple[str, str]] = [
-        ("CLICKUP", "tags"),
-        ("CONFLUENCE", "labels"),
-        ("DISCOURSE", "tags"),
-        ("FRESHDESK", "emails"),
-        ("GITHUB", "assignees"),
-        ("GITHUB", "labels"),
-        ("GURU", "tags"),
-        ("GURU", "folders"),
-        ("HUBSPOT", "associated_contact_ids"),
-        ("HUBSPOT", "associated_company_ids"),
-        ("HUBSPOT", "associated_deal_ids"),
-        ("HUBSPOT", "associated_ticket_ids"),
-        ("JIRA", "labels"),
-        ("MEDIAWIKI", "categories"),
-        ("ZENDESK", "labels"),
-        ("ZENDESK", "content_tags"),
-    ]
-
-    bind = op.get_bind()
-    for source, key in LIST_METADATA:
-        bind.execute(
-            sa.text(
-                f"""
-                UPDATE tag
-                SET is_list = true
-                WHERE tag_key = '{key}'
-                AND source = '{source}'
-                """
-            )
-        )
-
-
-def set_is_list_for_list_tags() -> None:
-    """
-    Sets is_list to true for all tags which have multiple values for a given
-    document, key, and source triplet. This only works if we remove old tags
-    from the database.
-    """
-    bind = op.get_bind()
-    bind.execute(
-        sa.text(
-            """
-            UPDATE tag
-            SET is_list = true
-            FROM (
-                SELECT DISTINCT tag.tag_key, tag.source
-                FROM tag
-                JOIN document__tag ON tag.id = document__tag.tag_id
-                GROUP BY tag.tag_key, tag.source, document__tag.document_id
-                HAVING count(*) > 1
-            ) AS list_tags
-            WHERE tag.tag_key = list_tags.tag_key
-            AND tag.source = list_tags.source
-            """
-        )
-    )
-
-
-def log_list_tags() -> None:
-    bind = op.get_bind()
-    result = bind.execute(
-        sa.text(
-            """
-            SELECT DISTINCT source, tag_key
-            FROM tag
-            WHERE is_list
-            ORDER BY source, tag_key
-            """
-        )
-    ).fetchall()
-    logger.info(
-        "List tags:\n" + "\n".join(f"  {source}: {key}" for source, key in result)
-    )
-
-
-def remove_old_tags() -> None:
-    """
-    Removes old tags from the database.
-    Previously, there was a bug where if a document got indexed with a tag and then
-    the document got reindexed, the old tag would not be removed.
-    This function removes those old tags by comparing it against the tags in vespa.
-    """
-    current_search_settings, future_search_settings = active_search_settings()
-    document_index = get_default_document_index(
-        current_search_settings, future_search_settings
-    )
-
-    # Get the index name
-    if hasattr(document_index, "index_name"):
-        index_name = document_index.index_name
-    else:
-        # Default index name if we can't get it from the document_index
-        index_name = "danswer_index"
-
-    for batch in _get_batch_documents_with_multiple_tags():
-        n_deleted = 0
-
-        for document_id in batch:
-            true_metadata = _get_vespa_metadata(document_id, index_name)
-            tags = _get_document_tags(document_id)
-
-            # identify document__tags to delete
-            to_delete: list[str] = []
-            for tag_id, tag_key, tag_value in tags:
-                true_val = true_metadata.get(tag_key, "")
-                if (isinstance(true_val, list) and tag_value not in true_val) or (
-                    isinstance(true_val, str) and tag_value != true_val
-                ):
-                    to_delete.append(str(tag_id))
-
-            if not to_delete:
-                continue
-
-            # delete old document__tags
-            bind = op.get_bind()
-            result = bind.execute(
-                sa.text(
-                    f"""
-                    DELETE FROM document__tag
-                    WHERE document_id = '{document_id}'
-                    AND tag_id IN ({','.join(to_delete)})
-                    """
-                )
-            )
-            n_deleted += result.rowcount
-        logger.info(f"Processed {len(batch)} documents and deleted {n_deleted} tags")
-
-
-def active_search_settings() -> tuple[SearchSettings, SearchSettings | None]:
-    result = op.get_bind().execute(
-        sa.text(
-            """
-        SELECT * FROM search_settings WHERE status = 'PRESENT' ORDER BY id DESC LIMIT 1
-        """
-        )
-    )
-    search_settings_fetch = result.fetchall()
-    search_settings = (
-        SearchSettings(**search_settings_fetch[0]._asdict())
-        if search_settings_fetch
-        else None
-    )
-
-    result2 = op.get_bind().execute(
-        sa.text(
-            """
-        SELECT * FROM search_settings WHERE status = 'FUTURE' ORDER BY id DESC LIMIT 1
-        """
-        )
-    )
-    search_settings_future_fetch = result2.fetchall()
-    search_settings_future = (
-        SearchSettings(**search_settings_future_fetch[0]._asdict())
-        if search_settings_future_fetch
-        else None
-    )
-
-    if not isinstance(search_settings, SearchSettings):
-        raise RuntimeError(
-            "current search settings is of type " + str(type(search_settings))
-        )
-    if (
-        not isinstance(search_settings_future, SearchSettings)
-        and search_settings_future is not None
-    ):
-        raise RuntimeError(
-            "future search settings is of type " + str(type(search_settings_future))
-        )
-
-    return search_settings, search_settings_future
-
-
-def _get_batch_documents_with_multiple_tags(
-    batch_size: int = 128,
-) -> Generator[list[str], None, None]:
-    """
-    Returns a list of document ids which contain a one to many tag.
-    The document may either contain a list metadata value, or may contain leftover
-    old tags from reindexing.
-    """
-    offset_clause = ""
-    bind = op.get_bind()
-
-    while True:
-        batch = bind.execute(
-            sa.text(
-                f"""
-                SELECT DISTINCT document__tag.document_id
-                FROM tag
-                JOIN document__tag ON tag.id = document__tag.tag_id
-                GROUP BY tag.tag_key, tag.source, document__tag.document_id
-                HAVING count(*) > 1 {offset_clause}
-                ORDER BY document__tag.document_id
-                LIMIT {batch_size}
-                """
-            )
-        ).fetchall()
-        if not batch:
-            break
-        doc_ids = [document_id for document_id, in batch]
-        yield doc_ids
-        offset_clause = f"AND document__tag.document_id > '{doc_ids[-1]}'"
-
-
-def _get_vespa_metadata(
-    document_id: str, index_name: str
-) -> dict[str, str | list[str]]:
-    url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)
-
-    # Document-Selector language
-    selection = (
-        f"{index_name}.document_id=='{document_id}' and {index_name}.chunk_id==0"
-    )
-
-    params: dict[str, str | int] = {
-        "selection": selection,
-        "wantedDocumentCount": 1,
-        "fieldSet": f"{index_name}:metadata",
-    }
-
-    with get_vespa_http_client() as client:
-        resp = client.get(url, params=params)
-        resp.raise_for_status()
-
-    docs = resp.json().get("documents", [])
-    if not docs:
-        raise RuntimeError(f"No chunk-0 found for document {document_id}")
-
-    # for some reason, metadata is a string
-    metadata = docs[0]["fields"]["metadata"]
-    return json.loads(metadata)
-
-
-def _get_document_tags(document_id: str) -> list[tuple[int, str, str]]:
-    bind = op.get_bind()
-    result = bind.execute(
-        sa.text(
-            f"""
-            SELECT tag.id, tag.tag_key, tag.tag_value
-            FROM tag
-            JOIN document__tag ON tag.id = document__tag.tag_id
-            WHERE document__tag.document_id = '{document_id}'
-            """
-        )
-    ).fetchall()
-    return cast(list[tuple[int, str, str]], result)
-
-
-def upgrade() -> None:
-    op.add_column(
-        "tag",
-        sa.Column("is_list", sa.Boolean(), nullable=False, server_default="false"),
-    )
-    op.drop_constraint(
-        constraint_name="_tag_key_value_source_uc",
-        table_name="tag",
-        type_="unique",
-    )
-    op.create_unique_constraint(
-        constraint_name="_tag_key_value_source_list_uc",
-        table_name="tag",
-        columns=["tag_key", "tag_value", "source", "is_list"],
-    )
-    set_is_list_for_known_tags()
-
-    if SKIP_TAG_FIX:
-        logger.warning(
-            "Skipping removal of old tags. "
-            "This can cause issues when using the knowledge graph, or "
-            "when filtering for documents by tags."
-        )
-        log_list_tags()
-        return
-
-    remove_old_tags()
-    set_is_list_for_list_tags()
-
-    # debug
-    log_list_tags()
-
-
-def downgrade() -> None:
-    # the migration adds and populates the is_list column, and removes old bugged tags
-    # there isn't a point in adding back the bugged tags, so we just drop the column
-    op.drop_constraint(
-        constraint_name="_tag_key_value_source_list_uc",
-        table_name="tag",
-        type_="unique",
-    )
-    op.create_unique_constraint(
-        constraint_name="_tag_key_value_source_uc",
-        table_name="tag",
-        columns=["tag_key", "tag_value", "source"],
-    )
-    op.drop_column("tag", "is_list")

backend/alembic/versions/b329d00a9ea6_adding_assistant_specific_user_.py

@@ -1,38 +0,0 @@
-"""Adding assistant-specific user preferences
-
-Revision ID: b329d00a9ea6
-Revises: f9b8c7d6e5a4
-Create Date: 2025-08-26 23:14:44.592985
-
-"""
-
-from alembic import op
-import fastapi_users_db_sqlalchemy
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-# revision identifiers, used by Alembic.
-revision = "b329d00a9ea6"
-down_revision = "f9b8c7d6e5a4"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    op.create_table(
-        "assistant__user_specific_config",
-        sa.Column("assistant_id", sa.Integer(), nullable=False),
-        sa.Column(
-            "user_id",
-            fastapi_users_db_sqlalchemy.generics.GUID(),
-            nullable=False,
-        ),
-        sa.Column("disabled_tool_ids", postgresql.ARRAY(sa.Integer()), nullable=False),
-        sa.ForeignKeyConstraint(["assistant_id"], ["persona.id"], ondelete="CASCADE"),
-        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
-        sa.PrimaryKeyConstraint("assistant_id", "user_id"),
-    )
-
-
-def downgrade() -> None:
-    op.drop_table("assistant__user_specific_config")

backend/alembic/versions/b558f51620b4_pause_finished_user_file_connectors.py

@@ -1,33 +0,0 @@
-"""Pause finished user file connectors
-
-Revision ID: b558f51620b4
-Revises: 90e3b9af7da4
-Create Date: 2025-08-15 17:17:02.456704
-
-"""
-
-from alembic import op
-
-
-# revision identifiers, used by Alembic.
-revision = "b558f51620b4"
-down_revision = "90e3b9af7da4"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Set all user file connector credential pairs with ACTIVE status to PAUSED
-    # This ensures user files don't continue to run indexing tasks after processing
-    op.execute(
-        """
-        UPDATE connector_credential_pair
-        SET status = 'PAUSED'
-        WHERE is_user_file = true
-        AND status = 'ACTIVE'
-        """
-    )
-
-
-def downgrade() -> None:
-    pass

backend/alembic/versions/bd7c3bf8beba_migrate_agent_responses_to_research_.py

@@ -1,147 +0,0 @@
-"""migrate_agent_sub_questions_to_research_iterations
-
-Revision ID: bd7c3bf8beba
-Revises: f8a9b2c3d4e5
-Create Date: 2025-08-18 11:33:27.098287
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "bd7c3bf8beba"
-down_revision = "f8a9b2c3d4e5"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Get connection to execute raw SQL
-    connection = op.get_bind()
-
-    # First, insert data into research_agent_iteration table
-    # This creates one iteration record per primary_question_id using the earliest time_created
-    connection.execute(
-        sa.text(
-            """
-            INSERT INTO research_agent_iteration (primary_question_id, created_at, iteration_nr, purpose, reasoning)
-            SELECT
-                primary_question_id,
-                MIN(time_created) as created_at,
-                1 as iteration_nr,
-                'Generating and researching subquestions' as purpose,
-                '(No previous reasoning)' as reasoning
-            FROM agent__sub_question
-            JOIN chat_message on agent__sub_question.primary_question_id = chat_message.id
-            WHERE primary_question_id IS NOT NULL
-                AND chat_message.is_agentic = true
-            GROUP BY primary_question_id
-            ON CONFLICT DO NOTHING;
-        """
-        )
-    )
-
-    # Then, insert data into research_agent_iteration_sub_step table
-    # This migrates each sub-question as a sub-step
-    connection.execute(
-        sa.text(
-            """
-            INSERT INTO research_agent_iteration_sub_step (
-                primary_question_id,
-                iteration_nr,
-                iteration_sub_step_nr,
-                created_at,
-                sub_step_instructions,
-                sub_step_tool_id,
-                sub_answer,
-                cited_doc_results
-            )
-            SELECT
-                primary_question_id,
-                1 as iteration_nr,
-                level_question_num as iteration_sub_step_nr,
-                time_created as created_at,
-                sub_question as sub_step_instructions,
-                1 as sub_step_tool_id,
-                sub_answer,
-                sub_question_doc_results as cited_doc_results
-            FROM agent__sub_question
-            JOIN chat_message on agent__sub_question.primary_question_id = chat_message.id
-            WHERE chat_message.is_agentic = true
-            AND primary_question_id IS NOT NULL
-            ON CONFLICT DO NOTHING;
-        """
-        )
-    )
-
-    # Update chat_message records: set legacy agentic type and answer purpose for existing agentic messages
-    connection.execute(
-        sa.text(
-            """
-            UPDATE chat_message
-            SET research_answer_purpose = 'ANSWER'
-            WHERE is_agentic = true
-            AND research_type IS NULL and
-                message_type = 'ASSISTANT';
-        """
-        )
-    )
-    connection.execute(
-        sa.text(
-            """
-            UPDATE chat_message
-            SET research_type = 'LEGACY_AGENTIC'
-            WHERE is_agentic = true
-            AND research_type IS NULL;
-        """
-        )
-    )
-
-
-def downgrade() -> None:
-    # Get connection to execute raw SQL
-    connection = op.get_bind()
-
-    # Note: This downgrade removes all research agent iteration data
-    # There's no way to perfectly restore the original agent__sub_question data
-    # if it was deleted after this migration
-
-    # Delete all research_agent_iteration_sub_step records that were migrated
-    connection.execute(
-        sa.text(
-            """
-            DELETE FROM research_agent_iteration_sub_step
-            USING chat_message
-            WHERE research_agent_iteration_sub_step.primary_question_id = chat_message.id
-            AND chat_message.research_type = 'LEGACY_AGENTIC';
-        """
-        )
-    )
-
-    # Delete all research_agent_iteration records that were migrated
-    connection.execute(
-        sa.text(
-            """
-            DELETE FROM research_agent_iteration
-            USING chat_message
-            WHERE research_agent_iteration.primary_question_id = chat_message.id
-            AND chat_message.research_type = 'LEGACY_AGENTIC';
-        """
-        )
-    )
-
-    # Revert chat_message updates: clear research fields for legacy agentic messages
-    connection.execute(
-        sa.text(
-            """
-            UPDATE chat_message
-            SET research_type = NULL,
-                research_answer_purpose = NULL
-            WHERE is_agentic = true
-            AND research_type = 'LEGACY_AGENTIC'
-            AND message_type = 'ASSISTANT';
-        """
-        )
-    )

backend/alembic/versions/c9e2cd766c29_add_s3_file_store_table.py

@@ -159,7 +159,7 @@ def _migrate_files_to_postgres() -> None:
 
     # only create external store if we have files to migrate. This line
     # makes it so we need to have S3/MinIO configured to run this migration.
-    external_store = get_s3_file_store()
+    external_store = get_s3_file_store(db_session=session)
 
     for i, file_id in enumerate(files_to_migrate, 1):
         print(f"Migrating file {i}/{total_files}: {file_id}")
@@ -219,7 +219,7 @@ def _migrate_files_to_external_storage() -> None:
     # Get database session
     bind = op.get_bind()
     session = Session(bind=bind)
-    external_store = get_s3_file_store()
+    external_store = get_s3_file_store(db_session=session)
 
     # Find all files currently stored in PostgreSQL (lobj_oid is not null)
     result = session.execute(
@@ -236,9 +236,6 @@ def _migrate_files_to_external_storage() -> None:
         print("No files found in PostgreSQL storage to migrate.")
         return
 
-    # might need to move this above the if statement when creating a new multi-tenant
-    # system. VERY extreme edge case.
-    external_store.initialize()
     print(f"Found {total_files} files to migrate from PostgreSQL to external storage.")
 
     _set_tenant_contextvar(session)

backend/alembic/versions/df0c7ad8a076_added_deletion_attempt_table.py

@@ -18,13 +18,11 @@ depends_on: None = None
 
 
 def upgrade() -> None:
-    op.execute("DROP TABLE IF EXISTS document CASCADE")
     op.create_table(
         "document",
         sa.Column("id", sa.String(), nullable=False),
         sa.PrimaryKeyConstraint("id"),
     )
-    op.execute("DROP TABLE IF EXISTS chunk CASCADE")
     op.create_table(
         "chunk",
         sa.Column("id", sa.String(), nullable=False),
@@ -45,7 +43,6 @@ def upgrade() -> None:
         ),
         sa.PrimaryKeyConstraint("id", "document_store_type"),
     )
-    op.execute("DROP TABLE IF EXISTS deletion_attempt CASCADE")
     op.create_table(
         "deletion_attempt",
         sa.Column("id", sa.Integer(), nullable=False),
@@ -87,7 +84,6 @@ def upgrade() -> None:
         ),
         sa.PrimaryKeyConstraint("id"),
     )
-    op.execute("DROP TABLE IF EXISTS document_by_connector_credential_pair CASCADE")
     op.create_table(
         "document_by_connector_credential_pair",
         sa.Column("id", sa.String(), nullable=False),
@@ -110,10 +106,7 @@ def upgrade() -> None:
 
 
 def downgrade() -> None:
-    # upstream tables first
     op.drop_table("document_by_connector_credential_pair")
     op.drop_table("deletion_attempt")
     op.drop_table("chunk")
-
-    # Alembic op.drop_table() has no "cascade" flag – issue raw SQL
-    op.execute("DROP TABLE IF EXISTS document CASCADE")
+    op.drop_table("document")

backend/alembic/versions/f8a9b2c3d4e5_add_research_answer_purpose_to_chat_message.py

@@ -1,30 +0,0 @@
-"""add research_answer_purpose to chat_message
-
-Revision ID: f8a9b2c3d4e5
-Revises: 5ae8240accb3
-Create Date: 2025-01-27 12:00:00.000000
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "f8a9b2c3d4e5"
-down_revision = "5ae8240accb3"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Add research_answer_purpose column to chat_message table
-    op.add_column(
-        "chat_message",
-        sa.Column("research_answer_purpose", sa.String(), nullable=True),
-    )
-
-
-def downgrade() -> None:
-    # Remove research_answer_purpose column from chat_message table
-    op.drop_column("chat_message", "research_answer_purpose")

backend/alembic/versions/f9b8c7d6e5a4_update_parent_question_id_foreign_key_to_research_agent_iteration.py

@@ -1,69 +0,0 @@
-"""remove foreign key constraints from research_agent_iteration_sub_step
-
-Revision ID: f9b8c7d6e5a4
-Revises: bd7c3bf8beba
-Create Date: 2025-01-27 12:00:00.000000
-
-"""
-
-from alembic import op
-import sqlalchemy as sa
-
-
-# revision identifiers, used by Alembic.
-revision = "f9b8c7d6e5a4"
-down_revision = "bd7c3bf8beba"
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # Drop the existing foreign key constraint for parent_question_id
-    op.drop_constraint(
-        "research_agent_iteration_sub_step_parent_question_id_fkey",
-        "research_agent_iteration_sub_step",
-        type_="foreignkey",
-    )
-
-    # Drop the parent_question_id column entirely
-    op.drop_column("research_agent_iteration_sub_step", "parent_question_id")
-
-    # Drop the foreign key constraint for primary_question_id to chat_message.id
-    # (keep the column as it's needed for the composite foreign key)
-    op.drop_constraint(
-        "research_agent_iteration_sub_step_primary_question_id_fkey",
-        "research_agent_iteration_sub_step",
-        type_="foreignkey",
-    )
-
-
-def downgrade() -> None:
-    # Restore the foreign key constraint for primary_question_id to chat_message.id
-    op.create_foreign_key(
-        "research_agent_iteration_sub_step_primary_question_id_fkey",
-        "research_agent_iteration_sub_step",
-        "chat_message",
-        ["primary_question_id"],
-        ["id"],
-        ondelete="CASCADE",
-    )
-
-    # Add back the parent_question_id column
-    op.add_column(
-        "research_agent_iteration_sub_step",
-        sa.Column(
-            "parent_question_id",
-            sa.Integer(),
-            nullable=True,
-        ),
-    )
-
-    # Restore the foreign key constraint pointing to research_agent_iteration_sub_step.id
-    op.create_foreign_key(
-        "research_agent_iteration_sub_step_parent_question_id_fkey",
-        "research_agent_iteration_sub_step",
-        "research_agent_iteration_sub_step",
-        ["parent_question_id"],
-        ["id"],
-        ondelete="CASCADE",
-    )

backend/ee/onyx/background/celery/apps/heavy.py

@@ -91,7 +91,7 @@ def export_query_history_task(
     with get_session_with_current_tenant() as db_session:
         try:
             stream.seek(0)
-            get_default_file_store().save_file(
+            get_default_file_store(db_session).save_file(
                 content=stream,
                 display_name=report_name,
                 file_origin=FileOrigin.QUERY_HISTORY_CSV,

backend/ee/onyx/background/celery/tasks/doc_permission_syncing/tasks.py

@@ -47,7 +47,6 @@ from onyx.connectors.factory import validate_ccpair_for_user
 from onyx.db.connector import mark_cc_pair_as_permissions_synced
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
 from onyx.db.document import get_document_ids_for_connector_credential_pair
-from onyx.db.document import get_documents_for_connector_credential_pair_limited_columns
 from onyx.db.document import upsert_document_by_connector_credential_pair
 from onyx.db.engine.sql_engine import get_session_with_current_tenant
 from onyx.db.engine.sql_engine import get_session_with_tenant
@@ -59,9 +58,7 @@ from onyx.db.models import ConnectorCredentialPair
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status
 from onyx.db.users import batch_add_ext_perm_user_if_not_exists
-from onyx.db.utils import DocumentRow
 from onyx.db.utils import is_retryable_sqlalchemy_error
-from onyx.db.utils import SortOrder
 from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
 from onyx.redis.redis_connector import RedisConnector
 from onyx.redis.redis_connector_doc_perm_sync import RedisConnectorPermissionSync
@@ -425,7 +422,7 @@ def connector_permission_sync_generator_task(
 
     lock: RedisLock = r.lock(
         OnyxRedisLocks.CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX
-        + f"_{redis_connector.cc_pair_id}",
+        + f"_{redis_connector.id}",
         timeout=CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT,
         thread_local=False,
     )
@@ -501,31 +498,16 @@ def connector_permission_sync_generator_task(
             # this is can be used to determine documents that are "missing" and thus
             # should no longer be accessible. The decision as to whether we should find
             # every document during the doc sync process is connector-specific.
-            def fetch_all_existing_docs_fn(
-                sort_order: SortOrder | None = None,
-            ) -> list[DocumentRow]:
-                result = get_documents_for_connector_credential_pair_limited_columns(
-                    db_session=db_session,
-                    connector_id=cc_pair.connector.id,
-                    credential_id=cc_pair.credential.id,
-                    sort_order=sort_order,
-                )
-                return list(result)
-
-            def fetch_all_existing_docs_ids_fn() -> list[str]:
-                result = get_document_ids_for_connector_credential_pair(
+            def fetch_all_existing_docs_fn() -> list[str]:
+                return get_document_ids_for_connector_credential_pair(
                     db_session=db_session,
                     connector_id=cc_pair.connector.id,
                     credential_id=cc_pair.credential.id,
                 )
-                return result
 
             doc_sync_func = sync_config.doc_sync_config.doc_sync_func
             document_external_accesses = doc_sync_func(
-                cc_pair,
-                fetch_all_existing_docs_fn,
-                fetch_all_existing_docs_ids_fn,
-                callback,
+                cc_pair, fetch_all_existing_docs_fn, callback
             )
 
             task_logger.info(

backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py

@@ -383,7 +383,7 @@ def connector_external_group_sync_generator_task(
 
     lock: RedisLock = r.lock(
         OnyxRedisLocks.CONNECTOR_EXTERNAL_GROUP_SYNC_LOCK_PREFIX
-        + f"_{redis_connector.cc_pair_id}",
+        + f"_{redis_connector.id}",
         timeout=CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT,
     )
 

backend/ee/onyx/chat/process_message.py

@@ -0,0 +1,38 @@
+from ee.onyx.server.query_and_chat.models import OneShotQAResponse
+from onyx.chat.models import AllCitations
+from onyx.chat.models import LLMRelevanceFilterResponse
+from onyx.chat.models import OnyxAnswerPiece
+from onyx.chat.models import QADocsResponse
+from onyx.chat.models import StreamingError
+from onyx.chat.process_message import ChatPacketStream
+from onyx.server.query_and_chat.models import ChatMessageDetail
+from onyx.utils.timing import log_function_time
+
+
+@log_function_time()
+def gather_stream_for_answer_api(
+    packets: ChatPacketStream,
+) -> OneShotQAResponse:
+    response = OneShotQAResponse()
+
+    answer = ""
+    for packet in packets:
+        if isinstance(packet, OnyxAnswerPiece) and packet.answer_piece:
+            answer += packet.answer_piece
+        elif isinstance(packet, QADocsResponse):
+            response.docs = packet
+            # Extraneous, provided for backwards compatibility
+            response.rephrase = packet.rephrased_query
+        elif isinstance(packet, StreamingError):
+            response.error_msg = packet.error
+        elif isinstance(packet, ChatMessageDetail):
+            response.chat_message_id = packet.message_id
+        elif isinstance(packet, LLMRelevanceFilterResponse):
+            response.llm_selected_doc_indices = packet.llm_selected_doc_indices
+        elif isinstance(packet, AllCitations):
+            response.citations = packet.citations
+
+    if answer:
+        response.answer = answer
+
+    return response

backend/ee/onyx/configs/app_configs.py

@@ -71,19 +71,6 @@ GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY = int(
 )
 
 
-#####
-# GitHub
-#####
-# In seconds, default is 5 minutes
-GITHUB_PERMISSION_DOC_SYNC_FREQUENCY = int(
-    os.environ.get("GITHUB_PERMISSION_DOC_SYNC_FREQUENCY") or 5 * 60
-)
-# In seconds, default is 5 minutes
-GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY = int(
-    os.environ.get("GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY") or 5 * 60
-)
-
-
 #####
 # Slack
 #####
@@ -102,19 +89,6 @@ TEAMS_PERMISSION_DOC_SYNC_FREQUENCY = int(
     os.environ.get("TEAMS_PERMISSION_DOC_SYNC_FREQUENCY") or 5 * 60
 )
 
-#####
-# SharePoint
-#####
-# In seconds, default is 30 minutes
-SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY = int(
-    os.environ.get("SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY") or 30 * 60
-)
-
-# In seconds, default is 5 minutes
-SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY = int(
-    os.environ.get("SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY") or 5 * 60
-)
-
 
 ####
 # Celery Job Frequency

backend/ee/onyx/db/usage_export.py

@@ -114,6 +114,7 @@ def get_all_usage_reports(db_session: Session) -> list[UsageReportMetadata]:
 
 
 def get_usage_report_data(
+    db_session: Session,
     report_display_name: str,
 ) -> IO:
     """
@@ -127,7 +128,7 @@ def get_usage_report_data(
     Returns:
         The usage report data.
     """
-    file_store = get_default_file_store()
+    file_store = get_default_file_store(db_session)
     # usage report may be very large, so don't load it all into memory
     return file_store.read_file(
         file_id=report_display_name, mode="b", use_tempfile=True

backend/ee/onyx/db/user_group.py

@@ -128,14 +128,11 @@ def validate_object_creation_for_user(
     target_group_ids: list[int] | None = None,
     object_is_public: bool | None = None,
     object_is_perm_sync: bool | None = None,
-    object_is_owned_by_user: bool = False,
-    object_is_new: bool = False,
 ) -> None:
     """
     All users can create/edit permission synced objects if they don't specify a group
     All admin actions are allowed.
-    Curators and global curators can create public objects.
-    Prevents other non-admins from creating/editing:
+    Prevents non-admins from creating/editing:
     - public objects
     - objects with no groups
     - objects that belong to a group they don't curate
@@ -146,23 +143,13 @@ def validate_object_creation_for_user(
     if not user or user.role == UserRole.ADMIN:
         return
 
-    # Allow curators and global curators to create public objects
-    # w/o associated groups IF the object is new/owned by them
-    if (
-        object_is_public
-        and user.role in [UserRole.CURATOR, UserRole.GLOBAL_CURATOR]
-        and (object_is_new or object_is_owned_by_user)
-    ):
-        return
-
-    if object_is_public and user.role == UserRole.BASIC:
-        detail = "User does not have permission to create public objects"
+    if object_is_public:
+        detail = "User does not have permission to create public credentials"
         logger.error(detail)
         raise HTTPException(
             status_code=400,
             detail=detail,
         )
-
     if not target_group_ids:
         detail = "Curators must specify 1+ groups"
         logger.error(detail)

backend/ee/onyx/document_index/vespa/app_config/cloud-services.xml.jinja

@@ -18,9 +18,9 @@
             <!-- <document type="danswer_chunk" mode="index" /> -->
 {{ document_elements }}
         </documents>
-        <nodes count="60">
-            <resources vcpu="8.0" memory="128.0Gb" architecture="arm64" storage-type="local"
-                disk="475.0Gb" />
+        <nodes count="75">
+            <resources vcpu="8.0" memory="64.0Gb" architecture="arm64" storage-type="local"
+                disk="474.0Gb" />
         </nodes>
         <engine>
             <proton>

backend/ee/onyx/external_permissions/confluence/doc_sync.py

@@ -6,7 +6,6 @@ https://confluence.atlassian.com/conf85/check-who-can-view-a-page-1283360557.htm
 from collections.abc import Generator
 
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from ee.onyx.external_permissions.utils import generic_doc_sync
 from onyx.access.models import DocExternalAccess
 from onyx.configs.constants import DocumentSource
@@ -26,7 +25,6 @@ CONFLUENCE_DOC_SYNC_LABEL = "confluence_doc_sync"
 def confluence_doc_sync(
     cc_pair: ConnectorCredentialPair,
     fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: IndexingHeartbeatInterface | None,
 ) -> Generator[DocExternalAccess, None, None]:
     """
@@ -45,7 +43,7 @@ def confluence_doc_sync(
 
     yield from generic_doc_sync(
         cc_pair=cc_pair,
-        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,
+        fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,
         callback=callback,
         doc_source=DocumentSource.CONFLUENCE,
         slim_connector=confluence_connector,

backend/ee/onyx/external_permissions/github/doc_sync.py

@@ -1,294 +0,0 @@
-import json
-from collections.abc import Generator
-
-from github import Github
-from github.Repository import Repository
-
-from ee.onyx.external_permissions.github.utils import fetch_repository_team_slugs
-from ee.onyx.external_permissions.github.utils import form_collaborators_group_id
-from ee.onyx.external_permissions.github.utils import form_organization_group_id
-from ee.onyx.external_permissions.github.utils import (
-    form_outside_collaborators_group_id,
-)
-from ee.onyx.external_permissions.github.utils import get_external_access_permission
-from ee.onyx.external_permissions.github.utils import get_repository_visibility
-from ee.onyx.external_permissions.github.utils import GitHubVisibility
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
-from onyx.access.models import DocExternalAccess
-from onyx.access.utils import build_ext_group_name_for_onyx
-from onyx.configs.constants import DocumentSource
-from onyx.connectors.github.connector import DocMetadata
-from onyx.connectors.github.connector import GithubConnector
-from onyx.db.models import ConnectorCredentialPair
-from onyx.db.utils import DocumentRow
-from onyx.db.utils import SortOrder
-from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-GITHUB_DOC_SYNC_LABEL = "github_doc_sync"
-
-
-def github_doc_sync(
-    cc_pair: ConnectorCredentialPair,
-    fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
-    callback: IndexingHeartbeatInterface | None = None,
-) -> Generator[DocExternalAccess, None, None]:
-    """
-    Sync GitHub documents with external access permissions.
-
-    This function checks each repository for visibility/team changes and updates
-    document permissions accordingly without using checkpoints.
-    """
-    logger.info(f"Starting GitHub document sync for CC pair ID: {cc_pair.id}")
-
-    # Initialize GitHub connector with credentials
-    github_connector: GithubConnector = GithubConnector(
-        **cc_pair.connector.connector_specific_config
-    )
-
-    github_connector.load_credentials(cc_pair.credential.credential_json)
-    logger.info("GitHub connector credentials loaded successfully")
-
-    if not github_connector.github_client:
-        logger.error("GitHub client initialization failed")
-        raise ValueError("github_client is required")
-
-    # Get all repositories from GitHub API
-    logger.info("Fetching all repositories from GitHub API")
-    try:
-        repos = []
-        if github_connector.repositories:
-            if "," in github_connector.repositories:
-                # Multiple repositories specified
-                repos = github_connector.get_github_repos(
-                    github_connector.github_client
-                )
-            else:
-                # Single repository
-                repos = [
-                    github_connector.get_github_repo(github_connector.github_client)
-                ]
-        else:
-            # All repositories
-            repos = github_connector.get_all_repos(github_connector.github_client)
-
-        logger.info(f"Found {len(repos)} repositories to check")
-    except Exception as e:
-        logger.error(f"Failed to fetch repositories: {e}")
-        raise
-
-    repo_to_doc_list_map: dict[str, list[DocumentRow]] = {}
-    # sort order is ascending because we want to get the oldest documents first
-    existing_docs: list[DocumentRow] = fetch_all_existing_docs_fn(
-        sort_order=SortOrder.ASC
-    )
-    logger.info(f"Found {len(existing_docs)} documents to check")
-    for doc in existing_docs:
-        try:
-            doc_metadata = DocMetadata.model_validate_json(json.dumps(doc.doc_metadata))
-            if doc_metadata.repo not in repo_to_doc_list_map:
-                repo_to_doc_list_map[doc_metadata.repo] = []
-            repo_to_doc_list_map[doc_metadata.repo].append(doc)
-        except Exception as e:
-            logger.error(f"Failed to parse doc metadata: {e} for doc {doc.id}")
-            continue
-    logger.info(f"Found {len(repo_to_doc_list_map)} documents to check")
-    # Process each repository individually
-    for repo in repos:
-        try:
-            logger.info(f"Processing repository: {repo.id} (name: {repo.name})")
-            repo_doc_list: list[DocumentRow] = repo_to_doc_list_map.get(
-                repo.full_name, []
-            )
-            if not repo_doc_list:
-                logger.warning(
-                    f"No documents found for repository {repo.id} ({repo.name})"
-                )
-                continue
-
-            current_external_group_ids = repo_doc_list[0].external_user_group_ids or []
-            # Check if repository has any permission changes
-            has_changes = _check_repository_for_changes(
-                repo=repo,
-                github_client=github_connector.github_client,
-                current_external_group_ids=current_external_group_ids,
-            )
-
-            if has_changes:
-                logger.info(
-                    f"Repository {repo.id} ({repo.name}) has changes, updating documents"
-                )
-
-                # Get new external access permissions for this repository
-                new_external_access = get_external_access_permission(
-                    repo, github_connector.github_client
-                )
-
-                logger.info(
-                    f"Found {len(repo_doc_list)} documents for repository {repo.full_name}"
-                )
-
-                # Yield updated external access for each document
-                for doc in repo_doc_list:
-                    if callback:
-                        callback.progress(GITHUB_DOC_SYNC_LABEL, 1)
-
-                    yield DocExternalAccess(
-                        doc_id=doc.id,
-                        external_access=new_external_access,
-                    )
-            else:
-                logger.info(
-                    f"Repository {repo.id} ({repo.name}) has no changes, skipping"
-                )
-        except Exception as e:
-            logger.error(f"Error processing repository {repo.id} ({repo.name}): {e}")
-
-    logger.info(f"GitHub document sync completed for CC pair ID: {cc_pair.id}")
-
-
-def _check_repository_for_changes(
-    repo: Repository,
-    github_client: Github,
-    current_external_group_ids: list[str],
-) -> bool:
-    """
-    Check if repository has any permission changes (visibility or team updates).
-    """
-    logger.info(f"Checking repository {repo.id} ({repo.name}) for changes")
-
-    # Check for repository visibility changes using the sample document data
-    if _is_repo_visibility_changed_from_groups(
-        repo=repo,
-        current_external_group_ids=current_external_group_ids,
-    ):
-        logger.info(f"Repository {repo.id} ({repo.name}) has visibility changes")
-        return True
-
-    # Check for team membership changes if repository is private
-    if get_repository_visibility(
-        repo
-    ) == GitHubVisibility.PRIVATE and _teams_updated_from_groups(
-        repo=repo,
-        github_client=github_client,
-        current_external_group_ids=current_external_group_ids,
-    ):
-        logger.info(f"Repository {repo.id} ({repo.name}) has team changes")
-        return True
-
-    logger.info(f"Repository {repo.id} ({repo.name}) has no changes")
-    return False
-
-
-def _is_repo_visibility_changed_from_groups(
-    repo: Repository,
-    current_external_group_ids: list[str],
-) -> bool:
-    """
-    Check if repository visibility has changed by analyzing existing external group IDs.
-
-    Args:
-        repo: GitHub repository object
-        current_external_group_ids: List of external group IDs from existing document
-
-    Returns:
-        True if visibility has changed
-    """
-    current_repo_visibility = get_repository_visibility(repo)
-    logger.info(f"Current repository visibility: {current_repo_visibility.value}")
-
-    # Build expected group IDs for current visibility
-    collaborators_group_id = build_ext_group_name_for_onyx(
-        source=DocumentSource.GITHUB,
-        ext_group_name=form_collaborators_group_id(repo.id),
-    )
-
-    org_group_id = None
-    if repo.organization:
-        org_group_id = build_ext_group_name_for_onyx(
-            source=DocumentSource.GITHUB,
-            ext_group_name=form_organization_group_id(repo.organization.id),
-        )
-
-    # Determine existing visibility from group IDs
-    has_collaborators_group = collaborators_group_id in current_external_group_ids
-    has_org_group = org_group_id and org_group_id in current_external_group_ids
-
-    if has_collaborators_group:
-        existing_repo_visibility = GitHubVisibility.PRIVATE
-    elif has_org_group:
-        existing_repo_visibility = GitHubVisibility.INTERNAL
-    else:
-        existing_repo_visibility = GitHubVisibility.PUBLIC
-
-    logger.info(f"Inferred existing visibility: {existing_repo_visibility.value}")
-
-    visibility_changed = existing_repo_visibility != current_repo_visibility
-    if visibility_changed:
-        logger.info(
-            f"Visibility changed for repo {repo.id} ({repo.name}): "
-            f"{existing_repo_visibility.value} -> {current_repo_visibility.value}"
-        )
-
-    return visibility_changed
-
-
-def _teams_updated_from_groups(
-    repo: Repository,
-    github_client: Github,
-    current_external_group_ids: list[str],
-) -> bool:
-    """
-    Check if repository team memberships have changed using existing group IDs.
-    """
-    # Fetch current team slugs for the repository
-    current_teams = fetch_repository_team_slugs(repo=repo, github_client=github_client)
-    logger.info(
-        f"Current teams for repository {repo.id} (name: {repo.name}): {current_teams}"
-    )
-
-    # Build group IDs to exclude from team comparison (non-team groups)
-    collaborators_group_id = build_ext_group_name_for_onyx(
-        source=DocumentSource.GITHUB,
-        ext_group_name=form_collaborators_group_id(repo.id),
-    )
-    outside_collaborators_group_id = build_ext_group_name_for_onyx(
-        source=DocumentSource.GITHUB,
-        ext_group_name=form_outside_collaborators_group_id(repo.id),
-    )
-    non_team_group_ids = {collaborators_group_id, outside_collaborators_group_id}
-
-    # Extract existing team IDs from current external group IDs
-    existing_team_ids = set()
-    for group_id in current_external_group_ids:
-        # Skip all non-team groups, keep only team groups
-        if group_id not in non_team_group_ids:
-            existing_team_ids.add(group_id)
-
-    # Note: existing_team_ids from DB are already prefixed (e.g., "github__team-slug")
-    # but current_teams from API are raw team slugs, so we need to add the prefix
-    current_team_ids = set()
-    for team_slug in current_teams:
-        team_group_id = build_ext_group_name_for_onyx(
-            source=DocumentSource.GITHUB,
-            ext_group_name=team_slug,
-        )
-        current_team_ids.add(team_group_id)
-
-    logger.info(
-        f"Existing team IDs: {existing_team_ids}, Current team IDs: {current_team_ids}"
-    )
-
-    # Compare actual team IDs to detect changes
-    teams_changed = current_team_ids != existing_team_ids
-    if teams_changed:
-        logger.info(
-            f"Team changes detected for repo {repo.id} (name: {repo.name}): "
-            f"existing={existing_team_ids}, current={current_team_ids}"
-        )
-
-    return teams_changed

backend/ee/onyx/external_permissions/github/group_sync.py

@@ -1,46 +0,0 @@
-from collections.abc import Generator
-
-from github import Repository
-
-from ee.onyx.db.external_perm import ExternalUserGroup
-from ee.onyx.external_permissions.github.utils import get_external_user_group
-from onyx.connectors.github.connector import GithubConnector
-from onyx.db.models import ConnectorCredentialPair
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def github_group_sync(
-    tenant_id: str,
-    cc_pair: ConnectorCredentialPair,
-) -> Generator[ExternalUserGroup, None, None]:
-    github_connector: GithubConnector = GithubConnector(
-        **cc_pair.connector.connector_specific_config
-    )
-    github_connector.load_credentials(cc_pair.credential.credential_json)
-    if not github_connector.github_client:
-        raise ValueError("github_client is required")
-
-    logger.info("Starting GitHub group sync...")
-    repos: list[Repository.Repository] = []
-    if github_connector.repositories:
-        if "," in github_connector.repositories:
-            # Multiple repositories specified
-            repos = github_connector.get_github_repos(github_connector.github_client)
-        else:
-            # Single repository (backward compatibility)
-            repos = [github_connector.get_github_repo(github_connector.github_client)]
-    else:
-        # All repositories
-        repos = github_connector.get_all_repos(github_connector.github_client)
-
-    for repo in repos:
-        try:
-            for external_group in get_external_user_group(
-                repo, github_connector.github_client
-            ):
-                logger.info(f"External group: {external_group}")
-                yield external_group
-        except Exception as e:
-            logger.error(f"Error processing repository {repo.id} ({repo.name}): {e}")

backend/ee/onyx/external_permissions/github/utils.py

@@ -1,488 +0,0 @@
-from collections.abc import Callable
-from enum import Enum
-from typing import List
-from typing import Optional
-from typing import Tuple
-from typing import TypeVar
-
-from github import Github
-from github import RateLimitExceededException
-from github.GithubException import GithubException
-from github.NamedUser import NamedUser
-from github.Organization import Organization
-from github.PaginatedList import PaginatedList
-from github.Repository import Repository
-from github.Team import Team
-from pydantic import BaseModel
-
-from ee.onyx.db.external_perm import ExternalUserGroup
-from onyx.access.models import ExternalAccess
-from onyx.access.utils import build_ext_group_name_for_onyx
-from onyx.configs.constants import DocumentSource
-from onyx.connectors.github.rate_limit_utils import sleep_after_rate_limit_exception
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-class GitHubVisibility(Enum):
-    """GitHub repository visibility options."""
-
-    PUBLIC = "public"
-    PRIVATE = "private"
-    INTERNAL = "internal"
-
-
-MAX_RETRY_COUNT = 3
-
-T = TypeVar("T")
-
-# Higher-order function to wrap GitHub operations with retry and exception handling
-
-
-def _run_with_retry(
-    operation: Callable[[], T],
-    description: str,
-    github_client: Github,
-    retry_count: int = 0,
-) -> Optional[T]:
-    """Execute a GitHub operation with retry on rate limit and exception handling."""
-    logger.debug(f"Starting operation '{description}', attempt {retry_count + 1}")
-    try:
-        result = operation()
-        logger.debug(f"Operation '{description}' completed successfully")
-        return result
-    except RateLimitExceededException:
-        if retry_count < MAX_RETRY_COUNT:
-            sleep_after_rate_limit_exception(github_client)
-            logger.warning(
-                f"Rate limit exceeded while {description}. Retrying... "
-                f"(attempt {retry_count + 1}/{MAX_RETRY_COUNT})"
-            )
-            return _run_with_retry(
-                operation, description, github_client, retry_count + 1
-            )
-        else:
-            error_msg = f"Max retries exceeded for {description}"
-            logger.exception(error_msg)
-            raise RuntimeError(error_msg)
-    except GithubException as e:
-        logger.warning(f"GitHub API error during {description}: {e}")
-        return None
-    except Exception as e:
-        logger.exception(f"Unexpected error during {description}: {e}")
-        return None
-
-
-class UserInfo(BaseModel):
-    """Represents a GitHub user with their basic information."""
-
-    login: str
-    name: Optional[str] = None
-    email: Optional[str] = None
-
-
-class TeamInfo(BaseModel):
-    """Represents a GitHub team with its members."""
-
-    name: str
-    slug: str
-    members: List[UserInfo]
-
-
-def _fetch_organization_members(
-    github_client: Github, org_name: str, retry_count: int = 0
-) -> List[UserInfo]:
-    """Fetch all organization members including owners and regular members."""
-    org_members: List[UserInfo] = []
-    logger.info(f"Fetching organization members for {org_name}")
-
-    org = _run_with_retry(
-        lambda: github_client.get_organization(org_name),
-        f"get organization {org_name}",
-        github_client,
-    )
-    if not org:
-        logger.error(f"Failed to fetch organization {org_name}")
-        raise RuntimeError(f"Failed to fetch organization {org_name}")
-
-    member_objs: PaginatedList[NamedUser] | list[NamedUser] = (
-        _run_with_retry(
-            lambda: org.get_members(filter_="all"),
-            f"get members for organization {org_name}",
-            github_client,
-        )
-        or []
-    )
-
-    for member in member_objs:
-        user_info = UserInfo(login=member.login, name=member.name, email=member.email)
-        org_members.append(user_info)
-
-    logger.info(f"Fetched {len(org_members)} members for organization {org_name}")
-    return org_members
-
-
-def _fetch_repository_teams_detailed(
-    repo: Repository, github_client: Github, retry_count: int = 0
-) -> List[TeamInfo]:
-    """Fetch teams with access to the repository and their members."""
-    teams_data: List[TeamInfo] = []
-    logger.info(f"Fetching teams for repository {repo.full_name}")
-
-    team_objs: PaginatedList[Team] | list[Team] = (
-        _run_with_retry(
-            lambda: repo.get_teams(),
-            f"get teams for repository {repo.full_name}",
-            github_client,
-        )
-        or []
-    )
-
-    for team in team_objs:
-        logger.info(
-            f"Processing team {team.name} (slug: {team.slug}) for repository {repo.full_name}"
-        )
-
-        members: PaginatedList[NamedUser] | list[NamedUser] = (
-            _run_with_retry(
-                lambda: team.get_members(),
-                f"get members for team {team.name}",
-                github_client,
-            )
-            or []
-        )
-
-        team_members = []
-        for m in members:
-            user_info = UserInfo(login=m.login, name=m.name, email=m.email)
-            team_members.append(user_info)
-
-        team_info = TeamInfo(name=team.name, slug=team.slug, members=team_members)
-        teams_data.append(team_info)
-        logger.info(f"Team {team.name} has {len(team_members)} members")
-
-    logger.info(f"Fetched {len(teams_data)} teams for repository {repo.full_name}")
-    return teams_data
-
-
-def fetch_repository_team_slugs(
-    repo: Repository, github_client: Github, retry_count: int = 0
-) -> List[str]:
-    """Fetch team slugs with access to the repository."""
-    logger.info(f"Fetching team slugs for repository {repo.full_name}")
-    teams_data: List[str] = []
-
-    team_objs: PaginatedList[Team] | list[Team] = (
-        _run_with_retry(
-            lambda: repo.get_teams(),
-            f"get teams for repository {repo.full_name}",
-            github_client,
-        )
-        or []
-    )
-
-    for team in team_objs:
-        teams_data.append(team.slug)
-
-    logger.info(f"Fetched {len(teams_data)} team slugs for repository {repo.full_name}")
-    return teams_data
-
-
-def _get_collaborators_and_outside_collaborators(
-    github_client: Github,
-    repo: Repository,
-) -> Tuple[List[UserInfo], List[UserInfo]]:
-    """Fetch and categorize collaborators into regular and outside collaborators."""
-    collaborators: List[UserInfo] = []
-    outside_collaborators: List[UserInfo] = []
-    logger.info(f"Fetching collaborators for repository {repo.full_name}")
-
-    repo_collaborators: PaginatedList[NamedUser] | list[NamedUser] = (
-        _run_with_retry(
-            lambda: repo.get_collaborators(),
-            f"get collaborators for repository {repo.full_name}",
-            github_client,
-        )
-        or []
-    )
-
-    for collaborator in repo_collaborators:
-        is_outside = False
-
-        # Check if collaborator is outside the organization
-        if repo.organization:
-            org: Organization | None = _run_with_retry(
-                lambda: github_client.get_organization(repo.organization.login),
-                f"get organization {repo.organization.login}",
-                github_client,
-            )
-
-            if org is not None:
-                org_obj = org
-                membership = _run_with_retry(
-                    lambda: org_obj.has_in_members(collaborator),
-                    f"check membership for {collaborator.login} in org {org_obj.login}",
-                    github_client,
-                )
-                is_outside = membership is not None and not membership
-
-        info = UserInfo(
-            login=collaborator.login, name=collaborator.name, email=collaborator.email
-        )
-        if repo.organization and is_outside:
-            outside_collaborators.append(info)
-        else:
-            collaborators.append(info)
-
-    logger.info(
-        f"Categorized {len(collaborators)} regular and {len(outside_collaborators)} outside collaborators for {repo.full_name}"
-    )
-    return collaborators, outside_collaborators
-
-
-def form_collaborators_group_id(repository_id: int) -> str:
-    """Generate group ID for repository collaborators."""
-    if not repository_id:
-        logger.exception("Repository ID is required to generate collaborators group ID")
-        raise ValueError("Repository ID must be set to generate group ID.")
-    group_id = f"{repository_id}_collaborators"
-    return group_id
-
-
-def form_organization_group_id(organization_id: int) -> str:
-    """Generate group ID for organization using organization ID."""
-    if not organization_id:
-        logger.exception(
-            "Organization ID is required to generate organization group ID"
-        )
-        raise ValueError("Organization ID must be set to generate group ID.")
-    group_id = f"{organization_id}_organization"
-    return group_id
-
-
-def form_outside_collaborators_group_id(repository_id: int) -> str:
-    """Generate group ID for outside collaborators."""
-    if not repository_id:
-        logger.exception(
-            "Repository ID is required to generate outside collaborators group ID"
-        )
-        raise ValueError("Repository ID must be set to generate group ID.")
-    group_id = f"{repository_id}_outside_collaborators"
-    return group_id
-
-
-def get_repository_visibility(repo: Repository) -> GitHubVisibility:
-    """
-    Get the visibility of a repository.
-    Returns GitHubVisibility enum member.
-    """
-    if hasattr(repo, "visibility"):
-        visibility = repo.visibility
-        logger.info(
-            f"Repository {repo.full_name} visibility from attribute: {visibility}"
-        )
-        try:
-            return GitHubVisibility(visibility)
-        except ValueError:
-            logger.warning(
-                f"Unknown visibility '{visibility}' for repo {repo.full_name}, defaulting to private"
-            )
-            return GitHubVisibility.PRIVATE
-
-    logger.info(f"Repository {repo.full_name} is private")
-    return GitHubVisibility.PRIVATE
-
-
-def get_external_access_permission(
-    repo: Repository, github_client: Github, add_prefix: bool = False
-) -> ExternalAccess:
-    """
-    Get the external access permission for a repository.
-    Uses group-based permissions for efficiency and scalability.
-
-    add_prefix: When this method is called during the initial permission sync via the connector,
-                the group ID isn't prefixed with the source while inserting the document record.
-                So in that case, set add_prefix to True, allowing the method itself to handle
-                prefixing. However, when the same method is invoked from doc_sync, our system
-                already adds the prefix to the group ID while processing the ExternalAccess object.
-    """
-    # We maintain collaborators, and outside collaborators as two separate groups
-    # instead of adding individual user emails to ExternalAccess.external_user_emails for two reasons:
-    # 1. Changes in repo collaborators (additions/removals) would require updating all documents.
-    # 2. Repo permissions can change without updating the repo's updated_at timestamp,
-    #    forcing full permission syncs for all documents every time, which is inefficient.
-
-    repo_visibility = get_repository_visibility(repo)
-    logger.info(
-        f"Generating ExternalAccess for {repo.full_name}: visibility={repo_visibility.value}"
-    )
-
-    if repo_visibility == GitHubVisibility.PUBLIC:
-        logger.info(
-            f"Repository {repo.full_name} is public - allowing access to all users"
-        )
-        return ExternalAccess(
-            external_user_emails=set(),
-            external_user_group_ids=set(),
-            is_public=True,
-        )
-    elif repo_visibility == GitHubVisibility.PRIVATE:
-        logger.info(
-            f"Repository {repo.full_name} is private - setting up restricted access"
-        )
-
-        collaborators_group_id = form_collaborators_group_id(repo.id)
-        outside_collaborators_group_id = form_outside_collaborators_group_id(repo.id)
-        if add_prefix:
-            collaborators_group_id = build_ext_group_name_for_onyx(
-                source=DocumentSource.GITHUB,
-                ext_group_name=collaborators_group_id,
-            )
-            outside_collaborators_group_id = build_ext_group_name_for_onyx(
-                source=DocumentSource.GITHUB,
-                ext_group_name=outside_collaborators_group_id,
-            )
-        group_ids = {collaborators_group_id, outside_collaborators_group_id}
-
-        team_slugs = fetch_repository_team_slugs(repo, github_client)
-        if add_prefix:
-            team_slugs = [
-                build_ext_group_name_for_onyx(
-                    source=DocumentSource.GITHUB,
-                    ext_group_name=slug,
-                )
-                for slug in team_slugs
-            ]
-        group_ids.update(team_slugs)
-
-        logger.info(f"ExternalAccess groups for {repo.full_name}: {group_ids}")
-        return ExternalAccess(
-            external_user_emails=set(),
-            external_user_group_ids=group_ids,
-            is_public=False,
-        )
-    else:
-        # Internal repositories - accessible to organization members
-        logger.info(
-            f"Repository {repo.full_name} is internal - accessible to org members"
-        )
-        org_group_id = form_organization_group_id(repo.organization.id)
-        if add_prefix:
-            org_group_id = build_ext_group_name_for_onyx(
-                source=DocumentSource.GITHUB,
-                ext_group_name=org_group_id,
-            )
-        group_ids = {org_group_id}
-        logger.info(f"ExternalAccess groups for {repo.full_name}: {group_ids}")
-        return ExternalAccess(
-            external_user_emails=set(),
-            external_user_group_ids=group_ids,
-            is_public=False,
-        )
-
-
-def get_external_user_group(
-    repo: Repository, github_client: Github
-) -> list[ExternalUserGroup]:
-    """
-    Get the external user group for a repository.
-    Creates ExternalUserGroup objects with actual user emails for each permission group.
-    """
-    repo_visibility = get_repository_visibility(repo)
-    logger.info(
-        f"Generating ExternalUserGroups for {repo.full_name}: visibility={repo_visibility.value}"
-    )
-
-    if repo_visibility == GitHubVisibility.PRIVATE:
-        logger.info(f"Processing private repository {repo.full_name}")
-
-        collaborators, outside_collaborators = (
-            _get_collaborators_and_outside_collaborators(github_client, repo)
-        )
-        teams = _fetch_repository_teams_detailed(repo, github_client)
-        external_user_groups = []
-
-        user_emails = set()
-        for collab in collaborators:
-            if collab.email:
-                user_emails.add(collab.email)
-            else:
-                logger.error(f"Collaborator {collab.login} has no email")
-
-        if user_emails:
-            collaborators_group = ExternalUserGroup(
-                id=form_collaborators_group_id(repo.id),
-                user_emails=list(user_emails),
-            )
-            external_user_groups.append(collaborators_group)
-            logger.info(f"Created collaborators group with {len(user_emails)} emails")
-
-        # Create group for outside collaborators
-        user_emails = set()
-        for collab in outside_collaborators:
-            if collab.email:
-                user_emails.add(collab.email)
-            else:
-                logger.error(f"Outside collaborator {collab.login} has no email")
-
-        if user_emails:
-            outside_collaborators_group = ExternalUserGroup(
-                id=form_outside_collaborators_group_id(repo.id),
-                user_emails=list(user_emails),
-            )
-            external_user_groups.append(outside_collaborators_group)
-            logger.info(
-                f"Created outside collaborators group with {len(user_emails)} emails"
-            )
-
-        # Create groups for teams
-        for team in teams:
-            user_emails = set()
-            for member in team.members:
-                if member.email:
-                    user_emails.add(member.email)
-                else:
-                    logger.error(f"Team member {member.login} has no email")
-
-            if user_emails:
-                team_group = ExternalUserGroup(
-                    id=team.slug,
-                    user_emails=list(user_emails),
-                )
-                external_user_groups.append(team_group)
-                logger.info(
-                    f"Created team group {team.name} with {len(user_emails)} emails"
-                )
-
-        logger.info(
-            f"Created {len(external_user_groups)} ExternalUserGroups for private repository {repo.full_name}"
-        )
-        return external_user_groups
-
-    if repo_visibility == GitHubVisibility.INTERNAL:
-        logger.info(f"Processing internal repository {repo.full_name}")
-
-        org_group_id = form_organization_group_id(repo.organization.id)
-        org_members = _fetch_organization_members(
-            github_client, repo.organization.login
-        )
-
-        user_emails = set()
-        for member in org_members:
-            if member.email:
-                user_emails.add(member.email)
-            else:
-                logger.error(f"Org member {member.login} has no email")
-
-        org_group = ExternalUserGroup(
-            id=org_group_id,
-            user_emails=list(user_emails),
-        )
-        logger.info(
-            f"Created organization group with {len(user_emails)} emails for internal repository {repo.full_name}"
-        )
-        return [org_group]
-
-    logger.info(f"Repository {repo.full_name} is public - no user groups needed")
-    return []

backend/ee/onyx/external_permissions/gmail/doc_sync.py

@@ -3,7 +3,6 @@ from datetime import datetime
 from datetime import timezone
 
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from onyx.access.models import DocExternalAccess
 from onyx.connectors.gmail.connector import GmailConnector
 from onyx.connectors.interfaces import GenerateSlimDocumentOutput
@@ -36,7 +35,6 @@ def _get_slim_doc_generator(
 def gmail_doc_sync(
     cc_pair: ConnectorCredentialPair,
     fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: IndexingHeartbeatInterface | None,
 ) -> Generator[DocExternalAccess, None, None]:
     """

backend/ee/onyx/external_permissions/google_drive/doc_sync.py

@@ -8,7 +8,6 @@ from ee.onyx.external_permissions.google_drive.permission_retrieval import (
     get_permissions_by_ids,
 )
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from onyx.access.models import DocExternalAccess
 from onyx.access.models import ExternalAccess
 from onyx.connectors.google_drive.connector import GoogleDriveConnector
@@ -41,28 +40,8 @@ def _get_slim_doc_generator(
     )
 
 
-def _merge_permissions_lists(
-    permission_lists: list[list[GoogleDrivePermission]],
-) -> list[GoogleDrivePermission]:
-    """
-    Merge a list of permission lists into a single list of permissions.
-    """
-    seen_permission_ids: set[str] = set()
-    merged_permissions: list[GoogleDrivePermission] = []
-    for permission_list in permission_lists:
-        for permission in permission_list:
-            if permission.id not in seen_permission_ids:
-                merged_permissions.append(permission)
-                seen_permission_ids.add(permission.id)
-
-    return merged_permissions
-
-
 def get_external_access_for_raw_gdrive_file(
-    file: GoogleDriveFileType,
-    company_domain: str,
-    retriever_drive_service: GoogleDriveService | None,
-    admin_drive_service: GoogleDriveService,
+    file: GoogleDriveFileType, company_domain: str, drive_service: GoogleDriveService
 ) -> ExternalAccess:
     """
     Get the external access for a raw Google Drive file.
@@ -83,28 +62,11 @@ def get_external_access_for_raw_gdrive_file(
             GoogleDrivePermission.from_drive_permission(p) for p in permissions
         ]
     elif permission_ids:
-
-        def _get_permissions(
-            drive_service: GoogleDriveService,
-        ) -> list[GoogleDrivePermission]:
-            return get_permissions_by_ids(
-                drive_service=drive_service,
-                doc_id=doc_id,
-                permission_ids=permission_ids,
-            )
-
-        permissions_list = _get_permissions(
-            retriever_drive_service or admin_drive_service
+        permissions_list = get_permissions_by_ids(
+            drive_service=drive_service,
+            doc_id=doc_id,
+            permission_ids=permission_ids,
         )
-        if len(permissions_list) != len(permission_ids) and retriever_drive_service:
-            logger.warning(
-                f"Failed to get all permissions for file {doc_id} with retriever service, "
-                "trying admin service"
-            )
-            backup_permissions_list = _get_permissions(admin_drive_service)
-            permissions_list = _merge_permissions_lists(
-                [permissions_list, backup_permissions_list]
-            )
 
     folder_ids_to_inherit_permissions_from: set[str] = set()
     user_emails: set[str] = set()
@@ -170,7 +132,6 @@ def get_external_access_for_raw_gdrive_file(
 def gdrive_doc_sync(
     cc_pair: ConnectorCredentialPair,
     fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: IndexingHeartbeatInterface | None,
 ) -> Generator[DocExternalAccess, None, None]:
     """

backend/ee/onyx/external_permissions/google_drive/group_sync.py

@@ -44,17 +44,11 @@ def _get_all_folders(
 
     TODO: tweak things so we can fetch deltas.
     """
-    MAX_FAILED_PERCENTAGE = 0.5
-
     all_folders: list[FolderInfo] = []
     seen_folder_ids: set[str] = set()
 
-    def _get_all_folders_for_user(
-        google_drive_connector: GoogleDriveConnector,
-        skip_folders_without_permissions: bool,
-        user_email: str,
-    ) -> None:
-        """Helper to get folders for a specific user + update shared seen_folder_ids"""
+    user_emails = google_drive_connector._get_all_user_emails()
+    for user_email in user_emails:
         drive_service = get_drive_service(
             google_drive_connector.creds,
             user_email,
@@ -104,20 +98,6 @@ def _get_all_folders(
                 )
             )
 
-    failed_count = 0
-    user_emails = google_drive_connector._get_all_user_emails()
-    for user_email in user_emails:
-        try:
-            _get_all_folders_for_user(
-                google_drive_connector, skip_folders_without_permissions, user_email
-            )
-        except Exception:
-            logger.exception(f"Error getting folders for user {user_email}")
-            failed_count += 1
-
-            if failed_count > MAX_FAILED_PERCENTAGE * len(user_emails):
-                raise RuntimeError("Too many failed folder fetches during group sync")
-
     return all_folders
 
 

backend/ee/onyx/external_permissions/jira/doc_sync.py

@@ -1,7 +1,6 @@
 from collections.abc import Generator
 
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from ee.onyx.external_permissions.utils import generic_doc_sync
 from onyx.access.models import DocExternalAccess
 from onyx.configs.constants import DocumentSource
@@ -18,7 +17,6 @@ JIRA_DOC_SYNC_TAG = "jira_doc_sync"
 def jira_doc_sync(
     cc_pair: ConnectorCredentialPair,
     fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: IndexingHeartbeatInterface | None = None,
 ) -> Generator[DocExternalAccess, None, None]:
     jira_connector = JiraConnector(
@@ -28,7 +26,7 @@ def jira_doc_sync(
 
     yield from generic_doc_sync(
         cc_pair=cc_pair,
-        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,
+        fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,
         callback=callback,
         doc_source=DocumentSource.JIRA,
         slim_connector=jira_connector,

backend/ee/onyx/external_permissions/perm_sync_types.py

@@ -2,45 +2,27 @@ from collections.abc import Callable
 from collections.abc import Generator
 from typing import Optional
 from typing import Protocol
+from typing import TYPE_CHECKING
 
-from ee.onyx.db.external_perm import ExternalUserGroup  # noqa
-from onyx.access.models import DocExternalAccess  # noqa
 from onyx.context.search.models import InferenceChunk
-from onyx.db.models import ConnectorCredentialPair  # noqa
-from onyx.db.utils import DocumentRow
-from onyx.db.utils import SortOrder
-from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface  # noqa
+
+# Avoid circular imports
+if TYPE_CHECKING:
+    from ee.onyx.db.external_perm import ExternalUserGroup  # noqa
+    from onyx.access.models import DocExternalAccess  # noqa
+    from onyx.db.models import ConnectorCredentialPair  # noqa
+    from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface  # noqa
 
 
 class FetchAllDocumentsFunction(Protocol):
-    """Protocol for a function that fetches documents for a connector credential pair.
+    """Protocol for a function that fetches all document IDs for a connector credential pair."""
 
-    This protocol defines the interface for functions that retrieve documents
-    from the database, typically used in permission synchronization workflows.
-    """
-
-    def __call__(
-        self,
-        sort_order: SortOrder | None,
-    ) -> list[DocumentRow]:
+    def __call__(self) -> list[str]:
         """
-        Fetches documents for a connector credential pair.
-        """
-        ...
+        Returns a list of document IDs for a connector credential pair.
 
-
-class FetchAllDocumentsIdsFunction(Protocol):
-    """Protocol for a function that fetches document IDs for a connector credential pair.
-
-    This protocol defines the interface for functions that retrieve document IDs
-    from the database, typically used in permission synchronization workflows.
-    """
-
-    def __call__(
-        self,
-    ) -> list[str]:
-        """
-        Fetches document IDs for a connector credential pair.
+        This is typically used to determine which documents should no longer be
+        accessible during the document sync process.
         """
         ...
 
@@ -48,20 +30,19 @@ class FetchAllDocumentsIdsFunction(Protocol):
 # Defining the input/output types for the sync functions
 DocSyncFuncType = Callable[
     [
-        ConnectorCredentialPair,
+        "ConnectorCredentialPair",
         FetchAllDocumentsFunction,
-        FetchAllDocumentsIdsFunction,
-        Optional[IndexingHeartbeatInterface],
+        Optional["IndexingHeartbeatInterface"],
     ],
-    Generator[DocExternalAccess, None, None],
+    Generator["DocExternalAccess", None, None],
 ]
 
 GroupSyncFuncType = Callable[
     [
         str,  # tenant_id
-        ConnectorCredentialPair,  # cc_pair
+        "ConnectorCredentialPair",  # cc_pair
     ],
-    Generator[ExternalUserGroup, None, None],
+    Generator["ExternalUserGroup", None, None],
 ]
 
 # list of chunks to be censored and the user email. returns censored chunks

backend/ee/onyx/external_permissions/sharepoint/doc_sync.py

@@ -1,36 +0,0 @@
-from collections.abc import Generator
-
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
-from ee.onyx.external_permissions.utils import generic_doc_sync
-from onyx.access.models import DocExternalAccess
-from onyx.configs.constants import DocumentSource
-from onyx.connectors.sharepoint.connector import SharepointConnector
-from onyx.db.models import ConnectorCredentialPair
-from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-SHAREPOINT_DOC_SYNC_TAG = "sharepoint_doc_sync"
-
-
-def sharepoint_doc_sync(
-    cc_pair: ConnectorCredentialPair,
-    fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
-    callback: IndexingHeartbeatInterface | None = None,
-) -> Generator[DocExternalAccess, None, None]:
-    sharepoint_connector = SharepointConnector(
-        **cc_pair.connector.connector_specific_config,
-    )
-    sharepoint_connector.load_credentials(cc_pair.credential.credential_json)
-
-    yield from generic_doc_sync(
-        cc_pair=cc_pair,
-        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,
-        callback=callback,
-        doc_source=DocumentSource.SHAREPOINT,
-        slim_connector=sharepoint_connector,
-        label=SHAREPOINT_DOC_SYNC_TAG,
-    )

backend/ee/onyx/external_permissions/sharepoint/group_sync.py

@@ -1,63 +0,0 @@
-from collections.abc import Generator
-
-from office365.sharepoint.client_context import ClientContext  # type: ignore[import-untyped]
-
-from ee.onyx.db.external_perm import ExternalUserGroup
-from ee.onyx.external_permissions.sharepoint.permission_utils import (
-    get_sharepoint_external_groups,
-)
-from onyx.connectors.sharepoint.connector import acquire_token_for_rest
-from onyx.connectors.sharepoint.connector import SharepointConnector
-from onyx.db.models import ConnectorCredentialPair
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def sharepoint_group_sync(
-    tenant_id: str,
-    cc_pair: ConnectorCredentialPair,
-) -> Generator[ExternalUserGroup, None, None]:
-    """Sync SharePoint groups and their members"""
-
-    # Get site URLs from connector config
-    connector_config = cc_pair.connector.connector_specific_config
-
-    # Create SharePoint connector instance and load credentials
-    connector = SharepointConnector(**connector_config)
-    connector.load_credentials(cc_pair.credential.credential_json)
-
-    if not connector.msal_app:
-        raise RuntimeError("MSAL app not initialized in connector")
-
-    if not connector.sp_tenant_domain:
-        raise RuntimeError("Tenant domain not initialized in connector")
-
-    # Get site descriptors from connector (either configured sites or all sites)
-    site_descriptors = connector.site_descriptors or connector.fetch_sites()
-
-    if not site_descriptors:
-        raise RuntimeError("No SharePoint sites found for group sync")
-
-    logger.info(f"Processing {len(site_descriptors)} sites for group sync")
-
-    msal_app = connector.msal_app
-    sp_tenant_domain = connector.sp_tenant_domain
-    # Process each site
-    for site_descriptor in site_descriptors:
-        logger.debug(f"Processing site: {site_descriptor.url}")
-
-        # Create client context for the site using connector's MSAL app
-        ctx = ClientContext(site_descriptor.url).with_access_token(
-            lambda: acquire_token_for_rest(msal_app, sp_tenant_domain)
-        )
-
-        # Get external groups for this site
-        external_groups = get_sharepoint_external_groups(ctx, connector.graph_client)
-
-        # Yield each group
-        for group in external_groups:
-            logger.debug(
-                f"Found group: {group.id} with {len(group.user_emails)} members"
-            )
-            yield group

backend/ee/onyx/external_permissions/sharepoint/permission_utils.py

@@ -1,684 +0,0 @@
-import re
-from collections import deque
-from typing import Any
-from urllib.parse import unquote
-from urllib.parse import urlparse
-
-from office365.graph_client import GraphClient  # type: ignore[import-untyped]
-from office365.onedrive.driveitems.driveItem import DriveItem  # type: ignore[import-untyped]
-from office365.runtime.client_request import ClientRequestException  # type: ignore
-from office365.sharepoint.client_context import ClientContext  # type: ignore[import-untyped]
-from office365.sharepoint.permissions.securable_object import RoleAssignmentCollection  # type: ignore[import-untyped]
-from pydantic import BaseModel
-
-from ee.onyx.db.external_perm import ExternalUserGroup
-from onyx.access.models import ExternalAccess
-from onyx.access.utils import build_ext_group_name_for_onyx
-from onyx.configs.constants import DocumentSource
-from onyx.connectors.sharepoint.connector import sleep_and_retry
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-# These values represent different types of SharePoint principals used in permission assignments
-USER_PRINCIPAL_TYPE = 1  # Individual user accounts
-ANONYMOUS_USER_PRINCIPAL_TYPE = 3  # Anonymous/unauthenticated users (public access)
-AZURE_AD_GROUP_PRINCIPAL_TYPE = 4  # Azure Active Directory security groups
-SHAREPOINT_GROUP_PRINCIPAL_TYPE = 8  # SharePoint site groups (local to the site)
-MICROSOFT_DOMAIN = ".onmicrosoft"
-# Limited Access role type, limited access is a travel through permission not a actual permission
-LIMITED_ACCESS_ROLE_TYPES = [1, 9]
-LIMITED_ACCESS_ROLE_NAMES = ["Limited Access", "Web-Only Limited Access"]
-
-
-class SharepointGroup(BaseModel):
-    model_config = {"frozen": True}
-
-    name: str
-    login_name: str
-    principal_type: int
-
-
-class GroupsResult(BaseModel):
-    groups_to_emails: dict[str, set[str]]
-    found_public_group: bool
-
-
-def _get_azuread_group_guid_by_name(
-    graph_client: GraphClient, group_name: str
-) -> str | None:
-    try:
-        # Search for groups by display name
-        groups = sleep_and_retry(
-            graph_client.groups.filter(f"displayName eq '{group_name}'").get(),
-            "get_azuread_group_guid_by_name",
-        )
-
-        if groups and len(groups) > 0:
-            return groups[0].id
-
-        return None
-
-    except Exception as e:
-        logger.error(f"Failed to get Azure AD group GUID for name {group_name}: {e}")
-        return None
-
-
-def _extract_guid_from_claims_token(claims_token: str) -> str | None:
-
-    try:
-        # Pattern to match GUID in claims token
-        # Claims tokens often have format: c:0o.c|provider|GUID_suffix
-        guid_pattern = r"([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})"
-
-        match = re.search(guid_pattern, claims_token, re.IGNORECASE)
-        if match:
-            return match.group(1)
-
-        return None
-
-    except Exception as e:
-        logger.error(f"Failed to extract GUID from claims token {claims_token}: {e}")
-        return None
-
-
-def _get_group_guid_from_identifier(
-    graph_client: GraphClient, identifier: str
-) -> str | None:
-    try:
-        # Check if it's already a GUID
-        guid_pattern = r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$"
-        if re.match(guid_pattern, identifier, re.IGNORECASE):
-            return identifier
-
-        # Check if it's a SharePoint claims token
-        if identifier.startswith("c:0") and "|" in identifier:
-            guid = _extract_guid_from_claims_token(identifier)
-            if guid:
-                logger.info(f"Extracted GUID {guid} from claims token {identifier}")
-                return guid
-
-        # Try to search by display name as fallback
-        return _get_azuread_group_guid_by_name(graph_client, identifier)
-
-    except Exception as e:
-        logger.error(f"Failed to get group GUID from identifier {identifier}: {e}")
-        return None
-
-
-def _get_security_group_owners(graph_client: GraphClient, group_id: str) -> list[str]:
-    try:
-        # Get group owners using Graph API
-        group = graph_client.groups[group_id]
-        owners = sleep_and_retry(
-            group.owners.get_all(page_loaded=lambda _: None),
-            "get_security_group_owners",
-        )
-
-        owner_emails: list[str] = []
-        logger.info(f"Owners: {owners}")
-
-        for owner in owners:
-            owner_data = owner.to_json()
-
-            # Extract email from the JSON data
-            mail: str | None = owner_data.get("mail")
-            user_principal_name: str | None = owner_data.get("userPrincipalName")
-
-            # Check if owner is a user and has an email
-            if mail:
-                if MICROSOFT_DOMAIN in mail:
-                    mail = mail.replace(MICROSOFT_DOMAIN, "")
-                owner_emails.append(mail)
-            elif user_principal_name:
-                if MICROSOFT_DOMAIN in user_principal_name:
-                    user_principal_name = user_principal_name.replace(
-                        MICROSOFT_DOMAIN, ""
-                    )
-                owner_emails.append(user_principal_name)
-
-        logger.info(
-            f"Retrieved {len(owner_emails)} owners from security group {group_id}"
-        )
-        return owner_emails
-
-    except Exception as e:
-        logger.error(f"Failed to get security group owners for group {group_id}: {e}")
-        return []
-
-
-def _get_sharepoint_list_item_id(drive_item: DriveItem) -> str | None:
-
-    try:
-        # First try to get the list item directly from the drive item
-        if hasattr(drive_item, "listItem"):
-            list_item = drive_item.listItem
-            if list_item:
-                # Load the list item properties to get the ID
-                sleep_and_retry(list_item.get(), "get_sharepoint_list_item_id")
-                if hasattr(list_item, "id") and list_item.id:
-                    return str(list_item.id)
-
-        # The SharePoint list item ID is typically available in the sharepointIds property
-        sharepoint_ids = getattr(drive_item, "sharepoint_ids", None)
-        if sharepoint_ids and hasattr(sharepoint_ids, "listItemId"):
-            return sharepoint_ids.listItemId
-
-        # Alternative: try to get it from the properties
-        properties = getattr(drive_item, "properties", None)
-        if properties:
-            # Sometimes the SharePoint list item ID is in the properties
-            for prop_name, prop_value in properties.items():
-                if "listitemid" in prop_name.lower():
-                    return str(prop_value)
-
-        return None
-    except Exception as e:
-        logger.error(
-            f"Error getting SharePoint list item ID for item {drive_item.id}: {e}"
-        )
-        raise e
-
-
-def _is_public_item(drive_item: DriveItem) -> bool:
-    is_public = False
-    try:
-        permissions = sleep_and_retry(
-            drive_item.permissions.get_all(page_loaded=lambda _: None), "is_public_item"
-        )
-        for permission in permissions:
-            if permission.link and (
-                permission.link.scope == "anonymous"
-                or permission.link.scope == "organization"
-            ):
-                is_public = True
-                break
-        return is_public
-    except Exception as e:
-        logger.error(f"Failed to check if item {drive_item.id} is public: {e}")
-        return False
-
-
-def _is_public_login_name(login_name: str) -> bool:
-    # Patterns that indicate public access
-    # This list is derived from the below link
-    # https://learn.microsoft.com/en-us/answers/questions/2085339/guid-in-the-loginname-of-site-user-everyone-except
-    public_login_patterns: list[str] = [
-        "c:0-.f|rolemanager|spo-grid-all-users/",
-        "c:0(.s|true",
-    ]
-    for pattern in public_login_patterns:
-        if pattern in login_name:
-            logger.info(f"Login name {login_name} is public")
-            return True
-    return False
-
-
-# AD groups allows same display name for multiple groups, so we need to add the GUID to the name
-def _get_group_name_with_suffix(
-    login_name: str, group_name: str, graph_client: GraphClient
-) -> str:
-    ad_group_suffix = _get_group_guid_from_identifier(graph_client, login_name)
-    return f"{group_name}_{ad_group_suffix}"
-
-
-def _get_sharepoint_groups(
-    client_context: ClientContext, group_name: str, graph_client: GraphClient
-) -> tuple[set[SharepointGroup], set[str]]:
-
-    groups: set[SharepointGroup] = set()
-    user_emails: set[str] = set()
-
-    def process_users(users: list[Any]) -> None:
-        nonlocal groups, user_emails
-
-        for user in users:
-            logger.debug(f"User: {user.to_json()}")
-            if user.principal_type == USER_PRINCIPAL_TYPE and hasattr(
-                user, "user_principal_name"
-            ):
-                if user.user_principal_name:
-                    email = user.user_principal_name
-                    if MICROSOFT_DOMAIN in email:
-                        email = email.replace(MICROSOFT_DOMAIN, "")
-                    user_emails.add(email)
-                else:
-                    logger.warning(
-                        f"User don't have a user principal name: {user.login_name}"
-                    )
-            elif user.principal_type in [
-                AZURE_AD_GROUP_PRINCIPAL_TYPE,
-                SHAREPOINT_GROUP_PRINCIPAL_TYPE,
-            ]:
-                name = user.title
-                if user.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:
-                    name = _get_group_name_with_suffix(
-                        user.login_name, name, graph_client
-                    )
-                groups.add(
-                    SharepointGroup(
-                        login_name=user.login_name,
-                        principal_type=user.principal_type,
-                        name=name,
-                    )
-                )
-
-    group = client_context.web.site_groups.get_by_name(group_name)
-    sleep_and_retry(
-        group.users.get_all(page_loaded=process_users), "get_sharepoint_groups"
-    )
-
-    return groups, user_emails
-
-
-def _get_azuread_groups(
-    graph_client: GraphClient, group_name: str
-) -> tuple[set[SharepointGroup], set[str]]:
-
-    group_id = _get_group_guid_from_identifier(graph_client, group_name)
-    if not group_id:
-        logger.error(f"Failed to get Azure AD group GUID for name {group_name}")
-        return set(), set()
-    group = graph_client.groups[group_id]
-    groups: set[SharepointGroup] = set()
-    user_emails: set[str] = set()
-
-    def process_members(members: list[Any]) -> None:
-        nonlocal groups, user_emails
-
-        for member in members:
-            member_data = member.to_json()
-            logger.debug(f"Member: {member_data}")
-            # Check for user-specific attributes
-            user_principal_name = member_data.get("userPrincipalName")
-            mail = member_data.get("mail")
-            display_name = member_data.get("displayName") or member_data.get(
-                "display_name"
-            )
-
-            # Check object attributes directly (if available)
-            is_user = False
-            is_group = False
-
-            # Users typically have userPrincipalName or mail
-            if user_principal_name or (mail and "@" in str(mail)):
-                is_user = True
-            # Groups typically have displayName but no userPrincipalName
-            elif display_name and not user_principal_name:
-                # Additional check: try to access group-specific properties
-                if (
-                    hasattr(member, "groupTypes")
-                    or member_data.get("groupTypes") is not None
-                ):
-                    is_group = True
-                # Or check if it has an 'id' field typical for groups
-                elif member_data.get("id") and not user_principal_name:
-                    is_group = True
-
-            # Check the object type name (fallback)
-            if not is_user and not is_group:
-                obj_type = type(member).__name__.lower()
-                if "user" in obj_type:
-                    is_user = True
-                elif "group" in obj_type:
-                    is_group = True
-
-            # Process based on identification
-            if is_user:
-                if user_principal_name:
-                    email = user_principal_name
-                    if MICROSOFT_DOMAIN in email:
-                        email = email.replace(MICROSOFT_DOMAIN, "")
-                    user_emails.add(email)
-                elif mail:
-                    email = mail
-                    if MICROSOFT_DOMAIN in email:
-                        email = email.replace(MICROSOFT_DOMAIN, "")
-                    user_emails.add(email)
-                logger.info(f"Added user: {user_principal_name or mail}")
-            elif is_group:
-                if not display_name:
-                    logger.error(f"No display name for group: {member_data.get('id')}")
-                    continue
-                name = _get_group_name_with_suffix(
-                    member_data.get("id", ""), display_name, graph_client
-                )
-                groups.add(
-                    SharepointGroup(
-                        login_name=member_data.get("id", ""),  # Use ID for groups
-                        principal_type=AZURE_AD_GROUP_PRINCIPAL_TYPE,
-                        name=name,
-                    )
-                )
-                logger.info(f"Added group: {name}")
-            else:
-                # Log unidentified members for debugging
-                logger.warning(f"Could not identify member type for: {member_data}")
-
-    sleep_and_retry(
-        group.members.get_all(page_loaded=process_members), "get_azuread_groups"
-    )
-
-    owner_emails = _get_security_group_owners(graph_client, group_id)
-    user_emails.update(owner_emails)
-
-    return groups, user_emails
-
-
-def _get_groups_and_members_recursively(
-    client_context: ClientContext,
-    graph_client: GraphClient,
-    groups: set[SharepointGroup],
-    is_group_sync: bool = False,
-) -> GroupsResult:
-    """
-    Get all groups and their members recursively.
-    """
-    group_queue: deque[SharepointGroup] = deque(groups)
-    visited_groups: set[str] = set()
-    visited_group_name_to_emails: dict[str, set[str]] = {}
-    found_public_group = False
-    while group_queue:
-        group = group_queue.popleft()
-        if group.login_name in visited_groups:
-            continue
-        visited_groups.add(group.login_name)
-        visited_group_name_to_emails[group.name] = set()
-        logger.info(
-            f"Processing group: {group.name} principal type: {group.principal_type}"
-        )
-        if group.principal_type == SHAREPOINT_GROUP_PRINCIPAL_TYPE:
-            group_info, user_emails = _get_sharepoint_groups(
-                client_context, group.login_name, graph_client
-            )
-            visited_group_name_to_emails[group.name].update(user_emails)
-            if group_info:
-                group_queue.extend(group_info)
-        if group.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:
-            try:
-                # if the site is public, we have default groups assigned to it, so we return early
-                if _is_public_login_name(group.login_name):
-                    found_public_group = True
-                    if not is_group_sync:
-                        return GroupsResult(
-                            groups_to_emails={}, found_public_group=True
-                        )
-                    else:
-                        # we don't want to sync public groups, so we skip them
-                        continue
-                group_info, user_emails = _get_azuread_groups(
-                    graph_client, group.login_name
-                )
-                visited_group_name_to_emails[group.name].update(user_emails)
-                if group_info:
-                    group_queue.extend(group_info)
-            except ClientRequestException as e:
-                # If the group is not found, we skip it. There is a chance that group is still referenced
-                # in sharepoint but it is removed from Azure AD. There is no actual documentation on this, but based on
-                # our testing we have seen this happen.
-                if e.response is not None and e.response.status_code == 404:
-                    logger.warning(f"Group {group.login_name} not found")
-                    continue
-                raise e
-
-    return GroupsResult(
-        groups_to_emails=visited_group_name_to_emails,
-        found_public_group=found_public_group,
-    )
-
-
-def get_external_access_from_sharepoint(
-    client_context: ClientContext,
-    graph_client: GraphClient,
-    drive_name: str | None,
-    drive_item: DriveItem | None,
-    site_page: dict[str, Any] | None,
-    add_prefix: bool = False,
-) -> ExternalAccess:
-    """
-    Get external access information from SharePoint.
-    """
-    groups: set[SharepointGroup] = set()
-    user_emails: set[str] = set()
-    group_ids: set[str] = set()
-
-    # Add all members to a processing set first
-    def add_user_and_group_to_sets(
-        role_assignments: RoleAssignmentCollection,
-    ) -> None:
-        nonlocal user_emails, groups
-        for assignment in role_assignments:
-            logger.debug(f"Assignment: {assignment.to_json()}")
-            if assignment.role_definition_bindings:
-                is_limited_access = True
-                for role_definition_binding in assignment.role_definition_bindings:
-                    if (
-                        role_definition_binding.role_type_kind
-                        not in LIMITED_ACCESS_ROLE_TYPES
-                        or role_definition_binding.name not in LIMITED_ACCESS_ROLE_NAMES
-                    ):
-                        is_limited_access = False
-                        break
-
-                # Skip if the role is only Limited Access, because this is not a actual permission its a travel through permission
-                if is_limited_access:
-                    logger.info(
-                        "Skipping assignment because it has only Limited Access role"
-                    )
-                    continue
-            if assignment.member:
-                member = assignment.member
-                if member.principal_type == USER_PRINCIPAL_TYPE and hasattr(
-                    member, "user_principal_name"
-                ):
-                    email = member.user_principal_name
-                    if MICROSOFT_DOMAIN in email:
-                        email = email.replace(MICROSOFT_DOMAIN, "")
-                    user_emails.add(email)
-                elif member.principal_type in [
-                    AZURE_AD_GROUP_PRINCIPAL_TYPE,
-                    SHAREPOINT_GROUP_PRINCIPAL_TYPE,
-                ]:
-                    name = member.title
-                    if member.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:
-                        name = _get_group_name_with_suffix(
-                            member.login_name, name, graph_client
-                        )
-                    groups.add(
-                        SharepointGroup(
-                            login_name=member.login_name,
-                            principal_type=member.principal_type,
-                            name=name,
-                        )
-                    )
-
-    if drive_item and drive_name:
-        # Here we check if the item have have any public links, if so we return early
-        is_public = _is_public_item(drive_item)
-        if is_public:
-            logger.info(f"Item {drive_item.id} is public")
-            return ExternalAccess(
-                external_user_emails=set(),
-                external_user_group_ids=set(),
-                is_public=True,
-            )
-
-        item_id = _get_sharepoint_list_item_id(drive_item)
-
-        if not item_id:
-            raise RuntimeError(
-                f"Failed to get SharePoint list item ID for item {drive_item.id}"
-            )
-
-        if drive_name == "Shared Documents":
-            drive_name = "Documents"
-
-        item = client_context.web.lists.get_by_title(drive_name).items.get_by_id(
-            item_id
-        )
-
-        sleep_and_retry(
-            item.role_assignments.expand(["Member", "RoleDefinitionBindings"]).get_all(
-                page_loaded=add_user_and_group_to_sets,
-            ),
-            "get_external_access_from_sharepoint",
-        )
-    elif site_page:
-        site_url = site_page.get("webUrl")
-        # Prefer server-relative URL to avoid OData filters that break on apostrophes
-        server_relative_url = unquote(urlparse(site_url).path)
-        file_obj = client_context.web.get_file_by_server_relative_url(
-            server_relative_url
-        )
-        item = file_obj.listItemAllFields
-
-        sleep_and_retry(
-            item.role_assignments.expand(["Member", "RoleDefinitionBindings"]).get_all(
-                page_loaded=add_user_and_group_to_sets,
-            ),
-            "get_external_access_from_sharepoint",
-        )
-    else:
-        raise RuntimeError("No drive item or site page provided")
-
-    groups_and_members: GroupsResult = _get_groups_and_members_recursively(
-        client_context, graph_client, groups
-    )
-
-    # If the site is public, w have default groups assigned to it, so we return early
-    if groups_and_members.found_public_group:
-        return ExternalAccess(
-            external_user_emails=set(),
-            external_user_group_ids=set(),
-            is_public=True,
-        )
-
-    for group_name, _ in groups_and_members.groups_to_emails.items():
-        if add_prefix:
-            group_name = build_ext_group_name_for_onyx(
-                group_name, DocumentSource.SHAREPOINT
-            )
-        group_ids.add(group_name.lower())
-
-    logger.info(f"User emails: {len(user_emails)}")
-    logger.info(f"Group IDs: {len(group_ids)}")
-
-    return ExternalAccess(
-        external_user_emails=user_emails,
-        external_user_group_ids=group_ids,
-        is_public=False,
-    )
-
-
-def get_sharepoint_external_groups(
-    client_context: ClientContext, graph_client: GraphClient
-) -> list[ExternalUserGroup]:
-
-    groups: set[SharepointGroup] = set()
-
-    def add_group_to_sets(role_assignments: RoleAssignmentCollection) -> None:
-        nonlocal groups
-        for assignment in role_assignments:
-            if assignment.role_definition_bindings:
-                is_limited_access = True
-                for role_definition_binding in assignment.role_definition_bindings:
-                    if (
-                        role_definition_binding.role_type_kind
-                        not in LIMITED_ACCESS_ROLE_TYPES
-                        or role_definition_binding.name not in LIMITED_ACCESS_ROLE_NAMES
-                    ):
-                        is_limited_access = False
-                        break
-
-                # Skip if the role assignment is only Limited Access, because this is not a actual permission its
-                #  a travel through permission
-                if is_limited_access:
-                    logger.info(
-                        "Skipping assignment because it has only Limited Access role"
-                    )
-                    continue
-            if assignment.member:
-                member = assignment.member
-                if member.principal_type in [
-                    AZURE_AD_GROUP_PRINCIPAL_TYPE,
-                    SHAREPOINT_GROUP_PRINCIPAL_TYPE,
-                ]:
-                    name = member.title
-                    if member.principal_type == AZURE_AD_GROUP_PRINCIPAL_TYPE:
-                        name = _get_group_name_with_suffix(
-                            member.login_name, name, graph_client
-                        )
-
-                    groups.add(
-                        SharepointGroup(
-                            login_name=member.login_name,
-                            principal_type=member.principal_type,
-                            name=name,
-                        )
-                    )
-
-    sleep_and_retry(
-        client_context.web.role_assignments.expand(
-            ["Member", "RoleDefinitionBindings"]
-        ).get_all(page_loaded=add_group_to_sets),
-        "get_sharepoint_external_groups",
-    )
-    groups_and_members: GroupsResult = _get_groups_and_members_recursively(
-        client_context, graph_client, groups, is_group_sync=True
-    )
-
-    # get all Azure AD groups because if any group is assigned to the drive item, we don't want to miss them
-    # We can't assign sharepoint groups to drive items or drives, so we don't need to get all sharepoint groups
-    azure_ad_groups = sleep_and_retry(
-        graph_client.groups.get_all(page_loaded=lambda _: None),
-        "get_sharepoint_external_groups:get_azure_ad_groups",
-    )
-    logger.info(f"Azure AD Groups: {len(azure_ad_groups)}")
-    identified_groups: set[str] = set(groups_and_members.groups_to_emails.keys())
-    ad_groups_to_emails: dict[str, set[str]] = {}
-    for group in azure_ad_groups:
-        # If the group is already identified, we don't need to get the members
-        if group.display_name in identified_groups:
-            continue
-        # AD groups allows same display name for multiple groups, so we need to add the GUID to the name
-        name = group.display_name
-        name = _get_group_name_with_suffix(group.id, name, graph_client)
-
-        members = sleep_and_retry(
-            group.members.get_all(page_loaded=lambda _: None),
-            "get_sharepoint_external_groups:get_azure_ad_groups:get_members",
-        )
-        for member in members:
-            member_data = member.to_json()
-            user_principal_name = member_data.get("userPrincipalName")
-            mail = member_data.get("mail")
-            if not ad_groups_to_emails.get(name):
-                ad_groups_to_emails[name] = set()
-            if user_principal_name:
-                if MICROSOFT_DOMAIN in user_principal_name:
-                    user_principal_name = user_principal_name.replace(
-                        MICROSOFT_DOMAIN, ""
-                    )
-                ad_groups_to_emails[name].add(user_principal_name)
-            elif mail:
-                if MICROSOFT_DOMAIN in mail:
-                    mail = mail.replace(MICROSOFT_DOMAIN, "")
-                ad_groups_to_emails[name].add(mail)
-
-    external_user_groups: list[ExternalUserGroup] = []
-    for group_name, emails in groups_and_members.groups_to_emails.items():
-        external_user_group = ExternalUserGroup(
-            id=group_name,
-            user_emails=list(emails),
-        )
-        external_user_groups.append(external_user_group)
-
-    for group_name, emails in ad_groups_to_emails.items():
-        external_user_group = ExternalUserGroup(
-            id=group_name,
-            user_emails=list(emails),
-        )
-        external_user_groups.append(external_user_group)
-
-    return external_user_groups

backend/ee/onyx/external_permissions/slack/doc_sync.py

@@ -3,7 +3,6 @@ from collections.abc import Generator
 from slack_sdk import WebClient
 
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from ee.onyx.external_permissions.slack.utils import fetch_user_id_to_email_map
 from onyx.access.models import DocExternalAccess
 from onyx.access.models import ExternalAccess
@@ -131,7 +130,6 @@ def _get_slack_document_access(
 def slack_doc_sync(
     cc_pair: ConnectorCredentialPair,
     fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: IndexingHeartbeatInterface | None,
 ) -> Generator[DocExternalAccess, None, None]:
     """

backend/ee/onyx/external_permissions/sync_params.py

@@ -7,18 +7,12 @@ from pydantic import BaseModel
 from ee.onyx.configs.app_configs import CONFLUENCE_PERMISSION_DOC_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import CONFLUENCE_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import DEFAULT_PERMISSION_DOC_SYNC_FREQUENCY
-from ee.onyx.configs.app_configs import GITHUB_PERMISSION_DOC_SYNC_FREQUENCY
-from ee.onyx.configs.app_configs import GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import GOOGLE_DRIVE_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import JIRA_PERMISSION_DOC_SYNC_FREQUENCY
-from ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY
-from ee.onyx.configs.app_configs import SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import SLACK_PERMISSION_DOC_SYNC_FREQUENCY
 from ee.onyx.configs.app_configs import TEAMS_PERMISSION_DOC_SYNC_FREQUENCY
 from ee.onyx.external_permissions.confluence.doc_sync import confluence_doc_sync
 from ee.onyx.external_permissions.confluence.group_sync import confluence_group_sync
-from ee.onyx.external_permissions.github.doc_sync import github_doc_sync
-from ee.onyx.external_permissions.github.group_sync import github_group_sync
 from ee.onyx.external_permissions.gmail.doc_sync import gmail_doc_sync
 from ee.onyx.external_permissions.google_drive.doc_sync import gdrive_doc_sync
 from ee.onyx.external_permissions.google_drive.group_sync import gdrive_group_sync
@@ -26,13 +20,10 @@ from ee.onyx.external_permissions.jira.doc_sync import jira_doc_sync
 from ee.onyx.external_permissions.perm_sync_types import CensoringFuncType
 from ee.onyx.external_permissions.perm_sync_types import DocSyncFuncType
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from ee.onyx.external_permissions.perm_sync_types import GroupSyncFuncType
 from ee.onyx.external_permissions.salesforce.postprocessing import (
     censor_salesforce_chunks,
 )
-from ee.onyx.external_permissions.sharepoint.doc_sync import sharepoint_doc_sync
-from ee.onyx.external_permissions.sharepoint.group_sync import sharepoint_group_sync
 from ee.onyx.external_permissions.slack.doc_sync import slack_doc_sync
 from ee.onyx.external_permissions.teams.doc_sync import teams_doc_sync
 from onyx.configs.constants import DocumentSource
@@ -72,7 +63,6 @@ class SyncConfig(BaseModel):
 def mock_doc_sync(
     cc_pair: "ConnectorCredentialPair",
     fetch_all_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: Optional["IndexingHeartbeatInterface"],
 ) -> Generator["DocExternalAccess", None, None]:
     """Mock doc sync function for testing - returns empty list since permissions are fetched during indexing"""
@@ -127,18 +117,6 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {
             initial_index_should_sync=False,
         ),
     ),
-    DocumentSource.GITHUB: SyncConfig(
-        doc_sync_config=DocSyncConfig(
-            doc_sync_frequency=GITHUB_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=github_doc_sync,
-            initial_index_should_sync=True,
-        ),
-        group_sync_config=GroupSyncConfig(
-            group_sync_frequency=GITHUB_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=github_group_sync,
-            group_sync_is_cc_pair_agnostic=False,
-        ),
-    ),
     DocumentSource.SALESFORCE: SyncConfig(
         censoring_config=CensoringConfig(
             chunk_censoring_func=censor_salesforce_chunks,
@@ -160,18 +138,6 @@ _SOURCE_TO_SYNC_CONFIG: dict[DocumentSource, SyncConfig] = {
             initial_index_should_sync=True,
         ),
     ),
-    DocumentSource.SHAREPOINT: SyncConfig(
-        doc_sync_config=DocSyncConfig(
-            doc_sync_frequency=SHAREPOINT_PERMISSION_DOC_SYNC_FREQUENCY,
-            doc_sync_func=sharepoint_doc_sync,
-            initial_index_should_sync=True,
-        ),
-        group_sync_config=GroupSyncConfig(
-            group_sync_frequency=SHAREPOINT_PERMISSION_GROUP_SYNC_FREQUENCY,
-            group_sync_func=sharepoint_group_sync,
-            group_sync_is_cc_pair_agnostic=False,
-        ),
-    ),
 }
 
 

backend/ee/onyx/external_permissions/teams/doc_sync.py

@@ -1,7 +1,6 @@
 from collections.abc import Generator
 
 from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
 from ee.onyx.external_permissions.utils import generic_doc_sync
 from onyx.access.models import DocExternalAccess
 from onyx.configs.constants import DocumentSource
@@ -19,7 +18,6 @@ TEAMS_DOC_SYNC_LABEL = "teams_doc_sync"
 def teams_doc_sync(
     cc_pair: ConnectorCredentialPair,
     fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     callback: IndexingHeartbeatInterface | None,
 ) -> Generator[DocExternalAccess, None, None]:
     teams_connector = TeamsConnector(
@@ -29,7 +27,7 @@ def teams_doc_sync(
 
     yield from generic_doc_sync(
         cc_pair=cc_pair,
-        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,
+        fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,
         callback=callback,
         doc_source=DocumentSource.TEAMS,
         slim_connector=teams_connector,

backend/ee/onyx/external_permissions/utils.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 
-from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsIdsFunction
+from ee.onyx.external_permissions.perm_sync_types import FetchAllDocumentsFunction
 from onyx.access.models import DocExternalAccess
 from onyx.access.models import ExternalAccess
 from onyx.configs.constants import DocumentSource
@@ -14,7 +14,7 @@ logger = setup_logger()
 
 def generic_doc_sync(
     cc_pair: ConnectorCredentialPair,
-    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
+    fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
     callback: IndexingHeartbeatInterface | None,
     doc_source: DocumentSource,
     slim_connector: SlimConnector,
@@ -62,9 +62,9 @@ def generic_doc_sync(
             )
 
     logger.info(f"Querying existing document IDs for CC Pair ID: {cc_pair.id=}")
-    existing_doc_ids: list[str] = fetch_all_existing_docs_ids_fn()
+    existing_doc_ids = set(fetch_all_existing_docs_fn())
 
-    missing_doc_ids = set(existing_doc_ids) - newly_fetched_doc_ids
+    missing_doc_ids = existing_doc_ids - newly_fetched_doc_ids
 
     if not missing_doc_ids:
         return

backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py

@@ -206,7 +206,7 @@ def _handle_standard_answers(
 
         restate_question_blocks = get_restate_blocks(
             msg=query_msg.message,
-            is_slash_command=message_info.is_slash_command,
+            is_bot_msg=message_info.is_bot_msg,
         )
 
         answer_blocks = build_standard_answer_blocks(

backend/ee/onyx/server/enterprise_settings/api.py

@@ -134,14 +134,15 @@ def ee_fetch_settings() -> EnterpriseSettings:
 def put_logo(
     file: UploadFile,
     is_logotype: bool = False,
+    db_session: Session = Depends(get_session),
     _: User | None = Depends(current_admin_user),
 ) -> None:
-    upload_logo(file=file, is_logotype=is_logotype)
+    upload_logo(file=file, db_session=db_session, is_logotype=is_logotype)
 
 
 def fetch_logo_helper(db_session: Session) -> Response:
     try:
-        file_store = get_default_file_store()
+        file_store = get_default_file_store(db_session)
         onyx_file = file_store.get_file_with_mime_type(get_logo_filename())
         if not onyx_file:
             raise ValueError("get_onyx_file returned None!")
@@ -157,7 +158,7 @@ def fetch_logo_helper(db_session: Session) -> Response:
 
 def fetch_logotype_helper(db_session: Session) -> Response:
     try:
-        file_store = get_default_file_store()
+        file_store = get_default_file_store(db_session)
         onyx_file = file_store.get_file_with_mime_type(get_logotype_filename())
         if not onyx_file:
             raise ValueError("get_onyx_file returned None!")

backend/ee/onyx/server/enterprise_settings/store.py

@@ -6,6 +6,7 @@ from typing import IO
 
 from fastapi import HTTPException
 from fastapi import UploadFile
+from sqlalchemy.orm import Session
 
 from ee.onyx.server.enterprise_settings.models import AnalyticsScriptUpload
 from ee.onyx.server.enterprise_settings.models import EnterpriseSettings
@@ -98,7 +99,9 @@ def guess_file_type(filename: str) -> str:
     return "application/octet-stream"
 
 
-def upload_logo(file: UploadFile | str, is_logotype: bool = False) -> bool:
+def upload_logo(
+    db_session: Session, file: UploadFile | str, is_logotype: bool = False
+) -> bool:
     content: IO[Any]
 
     if isinstance(file, str):
@@ -126,7 +129,7 @@ def upload_logo(file: UploadFile | str, is_logotype: bool = False) -> bool:
         display_name = file.filename
         file_type = file.content_type or "image/jpeg"
 
-    file_store = get_default_file_store()
+    file_store = get_default_file_store(db_session)
     file_store.save_file(
         content=content,
         display_name=display_name,

backend/ee/onyx/server/query_and_chat/chat_backend.py

@@ -1,22 +1,42 @@
+import re
+from typing import cast
+
 from fastapi import APIRouter
 from fastapi import Depends
 from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
+from ee.onyx.server.query_and_chat.models import AgentAnswer
+from ee.onyx.server.query_and_chat.models import AgentSubQuery
+from ee.onyx.server.query_and_chat.models import AgentSubQuestion
 from ee.onyx.server.query_and_chat.models import BasicCreateChatMessageRequest
 from ee.onyx.server.query_and_chat.models import (
     BasicCreateChatMessageWithHistoryRequest,
 )
+from ee.onyx.server.query_and_chat.models import ChatBasicResponse
 from onyx.auth.users import current_user
 from onyx.chat.chat_utils import combine_message_thread
 from onyx.chat.chat_utils import create_chat_chain
-from onyx.chat.models import ChatBasicResponse
-from onyx.chat.process_message import gather_stream
+from onyx.chat.models import AgentAnswerPiece
+from onyx.chat.models import AllCitations
+from onyx.chat.models import ExtendedToolResponse
+from onyx.chat.models import FinalUsedContextDocsResponse
+from onyx.chat.models import LlmDoc
+from onyx.chat.models import LLMRelevanceFilterResponse
+from onyx.chat.models import OnyxAnswerPiece
+from onyx.chat.models import QADocsResponse
+from onyx.chat.models import RefinedAnswerImprovement
+from onyx.chat.models import StreamingError
+from onyx.chat.models import SubQueryPiece
+from onyx.chat.models import SubQuestionIdentifier
+from onyx.chat.models import SubQuestionPiece
+from onyx.chat.process_message import ChatPacketStream
 from onyx.chat.process_message import stream_chat_message_objects
 from onyx.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from onyx.configs.constants import MessageType
 from onyx.context.search.models import OptionalSearchSetting
 from onyx.context.search.models import RetrievalDetails
+from onyx.context.search.models import SavedSearchDoc
 from onyx.db.chat import create_chat_session
 from onyx.db.chat import create_new_chat_message
 from onyx.db.chat import get_or_create_root_message
@@ -25,6 +45,7 @@ from onyx.db.models import User
 from onyx.llm.factory import get_llms_for_persona
 from onyx.natural_language_processing.utils import get_tokenizer
 from onyx.secondary_llm_flows.query_expansion import thread_based_query_rephrase
+from onyx.server.query_and_chat.models import ChatMessageDetail
 from onyx.server.query_and_chat.models import CreateChatMessageRequest
 from onyx.utils.logger import setup_logger
 
@@ -33,6 +54,177 @@ logger = setup_logger()
 router = APIRouter(prefix="/chat")
 
 
+def _get_final_context_doc_indices(
+    final_context_docs: list[LlmDoc] | None,
+    top_docs: list[SavedSearchDoc] | None,
+) -> list[int] | None:
+    """
+    this function returns a list of indices of the simple search docs
+    that were actually fed to the LLM.
+    """
+    if final_context_docs is None or top_docs is None:
+        return None
+
+    final_context_doc_ids = {doc.document_id for doc in final_context_docs}
+    return [
+        i for i, doc in enumerate(top_docs) if doc.document_id in final_context_doc_ids
+    ]
+
+
+def _convert_packet_stream_to_response(
+    packets: ChatPacketStream,
+) -> ChatBasicResponse:
+    response = ChatBasicResponse()
+    final_context_docs: list[LlmDoc] = []
+
+    answer = ""
+
+    # accumulate stream data with these dicts
+    agent_sub_questions: dict[tuple[int, int], AgentSubQuestion] = {}
+    agent_answers: dict[tuple[int, int], AgentAnswer] = {}
+    agent_sub_queries: dict[tuple[int, int, int], AgentSubQuery] = {}
+
+    for packet in packets:
+        if isinstance(packet, OnyxAnswerPiece) and packet.answer_piece:
+            answer += packet.answer_piece
+        elif isinstance(packet, QADocsResponse):
+            response.top_documents = packet.top_documents
+
+            # This is a no-op if agent_sub_questions hasn't already been filled
+            if packet.level is not None and packet.level_question_num is not None:
+                id = (packet.level, packet.level_question_num)
+                if id in agent_sub_questions:
+                    agent_sub_questions[id].document_ids = [
+                        saved_search_doc.document_id
+                        for saved_search_doc in packet.top_documents
+                    ]
+        elif isinstance(packet, StreamingError):
+            response.error_msg = packet.error
+        elif isinstance(packet, ChatMessageDetail):
+            response.message_id = packet.message_id
+        elif isinstance(packet, LLMRelevanceFilterResponse):
+            response.llm_selected_doc_indices = packet.llm_selected_doc_indices
+
+            # TODO: deprecate `llm_chunks_indices`
+            response.llm_chunks_indices = packet.llm_selected_doc_indices
+        elif isinstance(packet, FinalUsedContextDocsResponse):
+            final_context_docs = packet.final_context_docs
+        elif isinstance(packet, AllCitations):
+            response.cited_documents = {
+                citation.citation_num: citation.document_id
+                for citation in packet.citations
+            }
+        # agentic packets
+        elif isinstance(packet, SubQuestionPiece):
+            if packet.level is not None and packet.level_question_num is not None:
+                id = (packet.level, packet.level_question_num)
+                if agent_sub_questions.get(id) is None:
+                    agent_sub_questions[id] = AgentSubQuestion(
+                        level=packet.level,
+                        level_question_num=packet.level_question_num,
+                        sub_question=packet.sub_question,
+                        document_ids=[],
+                    )
+                else:
+                    agent_sub_questions[id].sub_question += packet.sub_question
+
+        elif isinstance(packet, AgentAnswerPiece):
+            if packet.level is not None and packet.level_question_num is not None:
+                id = (packet.level, packet.level_question_num)
+                if agent_answers.get(id) is None:
+                    agent_answers[id] = AgentAnswer(
+                        level=packet.level,
+                        level_question_num=packet.level_question_num,
+                        answer=packet.answer_piece,
+                        answer_type=packet.answer_type,
+                    )
+                else:
+                    agent_answers[id].answer += packet.answer_piece
+        elif isinstance(packet, SubQueryPiece):
+            if packet.level is not None and packet.level_question_num is not None:
+                sub_query_id = (
+                    packet.level,
+                    packet.level_question_num,
+                    packet.query_id,
+                )
+                if agent_sub_queries.get(sub_query_id) is None:
+                    agent_sub_queries[sub_query_id] = AgentSubQuery(
+                        level=packet.level,
+                        level_question_num=packet.level_question_num,
+                        sub_query=packet.sub_query,
+                        query_id=packet.query_id,
+                    )
+                else:
+                    agent_sub_queries[sub_query_id].sub_query += packet.sub_query
+        elif isinstance(packet, ExtendedToolResponse):
+            # we shouldn't get this ... it gets intercepted and translated to QADocsResponse
+            logger.warning(
+                "_convert_packet_stream_to_response: Unexpected chat packet type ExtendedToolResponse!"
+            )
+        elif isinstance(packet, RefinedAnswerImprovement):
+            response.agent_refined_answer_improvement = (
+                packet.refined_answer_improvement
+            )
+        else:
+            logger.warning(
+                f"_convert_packet_stream_to_response - Unrecognized chat packet: type={type(packet)}"
+            )
+
+    response.final_context_doc_indices = _get_final_context_doc_indices(
+        final_context_docs, response.top_documents
+    )
+
+    # organize / sort agent metadata for output
+    if len(agent_sub_questions) > 0:
+        response.agent_sub_questions = cast(
+            dict[int, list[AgentSubQuestion]],
+            SubQuestionIdentifier.make_dict_by_level(agent_sub_questions),
+        )
+
+    if len(agent_answers) > 0:
+        # return the agent_level_answer from the first level or the last one depending
+        # on agent_refined_answer_improvement
+        response.agent_answers = cast(
+            dict[int, list[AgentAnswer]],
+            SubQuestionIdentifier.make_dict_by_level(agent_answers),
+        )
+        if response.agent_answers:
+            selected_answer_level = (
+                0
+                if not response.agent_refined_answer_improvement
+                else len(response.agent_answers) - 1
+            )
+            level_answers = response.agent_answers[selected_answer_level]
+            for level_answer in level_answers:
+                if level_answer.answer_type != "agent_level_answer":
+                    continue
+
+                answer = level_answer.answer
+                break
+
+    if len(agent_sub_queries) > 0:
+        # subqueries are often emitted with trailing whitespace ... clean it up here
+        # perhaps fix at the source?
+        for v in agent_sub_queries.values():
+            v.sub_query = v.sub_query.strip()
+
+        response.agent_sub_queries = (
+            AgentSubQuery.make_dict_by_level_and_question_index(agent_sub_queries)
+        )
+
+    response.answer = answer
+    if answer:
+        response.answer_citationless = remove_answer_citations(answer)
+
+    return response
+
+
+def remove_answer_citations(answer: str) -> str:
+    pattern = r"\s*\[\[\d+\]\]\(http[s]?://[^\s]+\)"
+
+    return re.sub(pattern, "", answer)
+
+
 @router.post("/send-message-simple-api")
 def handle_simplified_chat_message(
     chat_message_req: BasicCreateChatMessageRequest,
@@ -45,36 +237,13 @@ def handle_simplified_chat_message(
     if not chat_message_req.message:
         raise HTTPException(status_code=400, detail="Empty chat message is invalid")
 
-    # Handle chat session creation if chat_session_id is not provided
-    if chat_message_req.chat_session_id is None:
-        if chat_message_req.persona_id is None:
-            raise HTTPException(
-                status_code=400,
-                detail="Either chat_session_id or persona_id must be provided",
-            )
-
-        # Create a new chat session with the provided persona_id
-        try:
-            new_chat_session = create_chat_session(
-                db_session=db_session,
-                description="",  # Leave empty for simple API
-                user_id=user.id if user else None,
-                persona_id=chat_message_req.persona_id,
-            )
-            chat_session_id = new_chat_session.id
-        except Exception as e:
-            logger.exception(e)
-            raise HTTPException(status_code=400, detail="Invalid Persona provided.")
-    else:
-        chat_session_id = chat_message_req.chat_session_id
-
     try:
         parent_message, _ = create_chat_chain(
-            chat_session_id=chat_session_id, db_session=db_session
+            chat_session_id=chat_message_req.chat_session_id, db_session=db_session
         )
     except Exception:
         parent_message = get_or_create_root_message(
-            chat_session_id=chat_session_id, db_session=db_session
+            chat_session_id=chat_message_req.chat_session_id, db_session=db_session
         )
 
     if (
@@ -89,7 +258,7 @@ def handle_simplified_chat_message(
         retrieval_options = chat_message_req.retrieval_options
 
     full_chat_msg_info = CreateChatMessageRequest(
-        chat_session_id=chat_session_id,
+        chat_session_id=chat_message_req.chat_session_id,
         parent_message_id=parent_message.id,
         message=chat_message_req.message,
         file_descriptors=[],
@@ -114,7 +283,7 @@ def handle_simplified_chat_message(
         enforce_chat_session_id_for_search_docs=False,
     )
 
-    return gather_stream(packets)
+    return _convert_packet_stream_to_response(packets)
 
 
 @router.post("/send-message-simple-with-history")
@@ -234,4 +403,4 @@ def handle_send_message_simple_with_history(
         enforce_chat_session_id_for_search_docs=False,
     )
 
-    return gather_stream(packets)
+    return _convert_packet_stream_to_response(packets)

backend/ee/onyx/server/query_and_chat/models.py

@@ -6,8 +6,10 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import model_validator
 
+from onyx.chat.models import CitationInfo
 from onyx.chat.models import PersonaOverrideConfig
 from onyx.chat.models import QADocsResponse
+from onyx.chat.models import SubQuestionIdentifier
 from onyx.chat.models import ThreadMessage
 from onyx.configs.constants import DocumentSource
 from onyx.context.search.enums import LLMEvaluationType
@@ -15,9 +17,8 @@ from onyx.context.search.enums import SearchType
 from onyx.context.search.models import ChunkContext
 from onyx.context.search.models import RerankingDetails
 from onyx.context.search.models import RetrievalDetails
+from onyx.context.search.models import SavedSearchDoc
 from onyx.server.manage.models import StandardAnswer
-from onyx.server.query_and_chat.streaming_models import CitationInfo
-from onyx.server.query_and_chat.streaming_models import SubQuestionIdentifier
 
 
 class StandardAnswerRequest(BaseModel):
@@ -40,13 +41,11 @@ class DocumentSearchRequest(ChunkContext):
 
 
 class BasicCreateChatMessageRequest(ChunkContext):
-    """If a chat_session_id is not provided, a persona_id must be provided to automatically create a new chat session
+    """Before creating messages, be sure to create a chat_session and get an id
     Note, for simplicity this option only allows for a single linear chain of messages
     """
 
-    chat_session_id: UUID | None = None
-    # Optional persona_id to create a new chat session if chat_session_id is not provided
-    persona_id: int | None = None
+    chat_session_id: UUID
     # New message contents
     message: str
     # Defaults to using retrieval with no additional filters
@@ -63,12 +62,6 @@ class BasicCreateChatMessageRequest(ChunkContext):
     # If True, uses agentic search instead of basic search
     use_agentic_search: bool = False
 
-    @model_validator(mode="after")
-    def validate_chat_session_or_persona(self) -> "BasicCreateChatMessageRequest":
-        if self.chat_session_id is None and self.persona_id is None:
-            raise ValueError("Either chat_session_id or persona_id must be provided")
-        return self
-
 
 class BasicCreateChatMessageWithHistoryRequest(ChunkContext):
     # Last element is the new query. All previous elements are historical context
@@ -155,6 +148,30 @@ class AgentSubQuery(SubQuestionIdentifier):
         return sorted_dict
 
 
+class ChatBasicResponse(BaseModel):
+    # This is built piece by piece, any of these can be None as the flow could break
+    answer: str | None = None
+    answer_citationless: str | None = None
+
+    top_documents: list[SavedSearchDoc] | None = None
+
+    error_msg: str | None = None
+    message_id: int | None = None
+    llm_selected_doc_indices: list[int] | None = None
+    final_context_doc_indices: list[int] | None = None
+    # this is a map of the citation number to the document id
+    cited_documents: dict[int, str] | None = None
+
+    # FOR BACKWARDS COMPATIBILITY
+    llm_chunks_indices: list[int] | None = None
+
+    # agentic fields
+    agent_sub_questions: dict[int, list[AgentSubQuestion]] | None = None
+    agent_answers: dict[int, list[AgentAnswer]] | None = None
+    agent_sub_queries: dict[int, dict[int, list[AgentSubQuery]]] | None = None
+    agent_refined_answer_improvement: bool | None = None
+
+
 class OneShotQARequest(ChunkContext):
     # Supports simplier APIs that don't deal with chat histories or message edits
     # Easier APIs to work with for developers
@@ -165,6 +182,7 @@ class OneShotQARequest(ChunkContext):
     prompt_id: int | None = None
     retrieval_options: RetrievalDetails = Field(default_factory=RetrievalDetails)
     rerank_settings: RerankingDetails | None = None
+    return_contexts: bool = False
 
     # allows the caller to specify the exact search query they want to use
     # can be used if the message sent to the LLM / query should not be the same
@@ -196,5 +214,6 @@ class OneShotQAResponse(BaseModel):
     rephrase: str | None = None
     citations: list[CitationInfo] | None = None
     docs: QADocsResponse | None = None
+    llm_selected_doc_indices: list[int] | None = None
     error_msg: str | None = None
     chat_message_id: int | None = None

backend/ee/onyx/server/query_and_chat/query_backend.py

@@ -8,6 +8,7 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from ee.onyx.chat.process_message import gather_stream_for_answer_api
 from ee.onyx.onyxbot.slack.handlers.handle_standard_answers import (
     oneoff_standard_answers,
 )
@@ -19,10 +20,8 @@ from ee.onyx.server.query_and_chat.models import StandardAnswerResponse
 from onyx.auth.users import current_user
 from onyx.chat.chat_utils import combine_message_thread
 from onyx.chat.chat_utils import prepare_chat_message_request
-from onyx.chat.models import AnswerStream
 from onyx.chat.models import PersonaOverrideConfig
-from onyx.chat.models import QADocsResponse
-from onyx.chat.process_message import gather_stream
+from onyx.chat.process_message import ChatPacketStream
 from onyx.chat.process_message import stream_chat_message_objects
 from onyx.configs.onyxbot_configs import MAX_THREAD_CONTEXT_PERCENTAGE
 from onyx.context.search.models import SavedSearchDocWithContent
@@ -40,7 +39,6 @@ from onyx.llm.factory import get_default_llms
 from onyx.llm.factory import get_llms_for_persona
 from onyx.llm.factory import get_main_llm_from_tuple
 from onyx.natural_language_processing.utils import get_tokenizer
-from onyx.server.query_and_chat.streaming_models import CitationInfo
 from onyx.server.utils import get_json_line
 from onyx.utils.logger import setup_logger
 
@@ -142,7 +140,7 @@ def get_answer_stream(
     query_request: OneShotQARequest,
     user: User | None = Depends(current_user),
     db_session: Session = Depends(get_session),
-) -> AnswerStream:
+) -> ChatPacketStream:
     query = query_request.messages[0].message
     logger.notice(f"Received query for Answer API: {query}")
 
@@ -207,6 +205,7 @@ def get_answer_stream(
         new_msg_req=request,
         user=user,
         db_session=db_session,
+        include_contexts=query_request.return_contexts,
     )
 
     return packets
@@ -220,28 +219,12 @@ def get_answer_with_citation(
 ) -> OneShotQAResponse:
     try:
         packets = get_answer_stream(request, user, db_session)
-        answer = gather_stream(packets)
+        answer = gather_stream_for_answer_api(packets)
 
         if answer.error_msg:
             raise RuntimeError(answer.error_msg)
 
-        return OneShotQAResponse(
-            answer=answer.answer,
-            chat_message_id=answer.message_id,
-            error_msg=answer.error_msg,
-            citations=[
-                CitationInfo(citation_num=i, document_id=doc_id)
-                for i, doc_id in answer.cited_documents.items()
-            ],
-            docs=QADocsResponse(
-                top_documents=answer.top_documents,
-                predicted_flow=None,
-                predicted_search=None,
-                applied_source_filters=None,
-                applied_time_cutoff=None,
-                recency_bias_multiplier=0.0,
-            ),
-        )
+        return answer
     except Exception as e:
         logger.error(f"Error in get_answer_with_citation: {str(e)}", exc_info=True)
         raise HTTPException(status_code=500, detail="An internal server error occurred")

backend/ee/onyx/server/query_history/api.py

@@ -358,7 +358,7 @@ def get_query_history_export_status(
     # If task is None, then it's possible that the task has already finished processing.
     # Therefore, we should then check if the export file has already been stored inside of the file-store.
     # If that *also* doesn't exist, then we can return a 404.
-    file_store = get_default_file_store()
+    file_store = get_default_file_store(db_session)
 
     report_name = construct_query_history_report_name(request_id)
     has_file = file_store.has_file(
@@ -385,7 +385,7 @@ def download_query_history_csv(
     ensure_query_history_is_enabled(disallowed=[QueryHistoryType.DISABLED])
 
     report_name = construct_query_history_report_name(request_id)
-    file_store = get_default_file_store()
+    file_store = get_default_file_store(db_session)
     has_file = file_store.has_file(
         file_id=report_name,
         file_origin=FileOrigin.QUERY_HISTORY_CSV,

backend/ee/onyx/server/reporting/usage_export_api.py

@@ -53,7 +53,7 @@ def read_usage_report(
     db_session: Session = Depends(get_session),
 ) -> Response:
     try:
-        file = get_usage_report_data(report_name)
+        file = get_usage_report_data(db_session, report_name)
     except ValueError as e:
         raise HTTPException(status_code=404, detail=str(e))
 

backend/ee/onyx/server/reporting/usage_export_generation.py

@@ -67,7 +67,7 @@ def generate_chat_messages_report(
         file_id = file_store.save_file(
             content=temp_file,
             display_name=file_name,
-            file_origin=FileOrigin.GENERATED_REPORT,
+            file_origin=FileOrigin.OTHER,
             file_type="text/csv",
         )
 
@@ -99,7 +99,7 @@ def generate_user_report(
         file_id = file_store.save_file(
             content=temp_file,
             display_name=file_name,
-            file_origin=FileOrigin.GENERATED_REPORT,
+            file_origin=FileOrigin.OTHER,
             file_type="text/csv",
         )
 
@@ -112,7 +112,7 @@ def create_new_usage_report(
     period: tuple[datetime, datetime] | None,
 ) -> UsageReportMetadata:
     report_id = str(uuid.uuid4())
-    file_store = get_default_file_store()
+    file_store = get_default_file_store(db_session)
 
     messages_file_id = generate_chat_messages_report(
         db_session, file_store, report_id, period

backend/ee/onyx/server/seeding.py

@@ -200,10 +200,10 @@ def _seed_enterprise_settings(seed_config: SeedConfiguration) -> None:
         store_ee_settings(final_enterprise_settings)
 
 
-def _seed_logo(logo_path: str | None) -> None:
+def _seed_logo(db_session: Session, logo_path: str | None) -> None:
     if logo_path:
         logger.notice("Uploading logo")
-        upload_logo(file=logo_path)
+        upload_logo(db_session=db_session, file=logo_path)
 
 
 def _seed_analytics_script(seed_config: SeedConfiguration) -> None:
@@ -245,7 +245,7 @@ def seed_db() -> None:
         if seed_config.custom_tools is not None:
             _seed_custom_tools(db_session, seed_config.custom_tools)
 
-        _seed_logo(seed_config.seeded_logo_path)
+        _seed_logo(db_session, seed_config.seeded_logo_path)
         _seed_enterprise_settings(seed_config)
         _seed_analytics_script(seed_config)
 

backend/ee/onyx/server/tenants/billing_api.py

@@ -10,12 +10,10 @@ from ee.onyx.server.tenants.billing import fetch_billing_information
 from ee.onyx.server.tenants.billing import fetch_stripe_checkout_session
 from ee.onyx.server.tenants.billing import fetch_tenant_stripe_information
 from ee.onyx.server.tenants.models import BillingInformation
-from ee.onyx.server.tenants.models import ProductGatingFullSyncRequest
 from ee.onyx.server.tenants.models import ProductGatingRequest
 from ee.onyx.server.tenants.models import ProductGatingResponse
 from ee.onyx.server.tenants.models import SubscriptionSessionResponse
 from ee.onyx.server.tenants.models import SubscriptionStatusResponse
-from ee.onyx.server.tenants.product_gating import overwrite_full_gated_set
 from ee.onyx.server.tenants.product_gating import store_product_gating
 from onyx.auth.users import User
 from onyx.configs.app_configs import WEB_DOMAIN
@@ -49,26 +47,6 @@ def gate_product(
         return ProductGatingResponse(updated=False, error=str(e))
 
 
-@router.post("/product-gating/full-sync")
-def gate_product_full_sync(
-    product_gating_request: ProductGatingFullSyncRequest,
-    _: None = Depends(control_plane_dep),
-) -> ProductGatingResponse:
-    """
-    Bulk operation to overwrite the entire gated tenant set.
-    This replaces all currently gated tenants with the provided list.
-    Gated tenants are not available to access the product and will be
-    directed to the billing page when their subscription has ended.
-    """
-    try:
-        overwrite_full_gated_set(product_gating_request.gated_tenant_ids)
-        return ProductGatingResponse(updated=True, error=None)
-
-    except Exception as e:
-        logger.exception("Failed to gate products during full sync")
-        return ProductGatingResponse(updated=False, error=str(e))
-
-
 @router.get("/billing-information")
 async def billing_information(
     _: User = Depends(current_admin_user),

backend/ee/onyx/server/tenants/models.py

@@ -19,10 +19,6 @@ class ProductGatingRequest(BaseModel):
     application_status: ApplicationStatus
 
 
-class ProductGatingFullSyncRequest(BaseModel):
-    gated_tenant_ids: list[str]
-
-
 class SubscriptionStatusResponse(BaseModel):
     subscribed: bool
 

backend/ee/onyx/server/tenants/product_gating.py

@@ -16,6 +16,10 @@ logger = setup_logger()
 def update_tenant_gating(tenant_id: str, status: ApplicationStatus) -> None:
     redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
 
+    # Store the full status
+    status_key = f"tenant:{tenant_id}:status"
+    redis_client.set(status_key, status.value)
+
     # Maintain the GATED_ACCESS set
     if status == ApplicationStatus.GATED_ACCESS:
         redis_client.sadd(GATED_TENANTS_KEY, tenant_id)
@@ -42,25 +46,6 @@ def store_product_gating(tenant_id: str, application_status: ApplicationStatus)
         raise
 
 
-def overwrite_full_gated_set(tenant_ids: list[str]) -> None:
-    redis_client = get_redis_client(tenant_id=ONYX_CLOUD_TENANT_ID)
-
-    pipeline = redis_client.pipeline()
-
-    # using pipeline doesn't automatically add the tenant_id prefix
-    full_gated_set_key = f"{ONYX_CLOUD_TENANT_ID}:{GATED_TENANTS_KEY}"
-
-    # Clear the existing set
-    pipeline.delete(full_gated_set_key)
-
-    # Add all tenant IDs to the set and set their status
-    for tenant_id in tenant_ids:
-        pipeline.sadd(full_gated_set_key, tenant_id)
-
-    # Execute all commands at once
-    pipeline.execute()
-
-
 def get_gated_tenants() -> set[str]:
     redis_client = get_redis_replica_client(tenant_id=ONYX_CLOUD_TENANT_ID)
     gated_tenants_bytes = cast(set[bytes], redis_client.smembers(GATED_TENANTS_KEY))

backend/model_server/constants.py

@@ -1,5 +1,34 @@
+from shared_configs.enums import EmbeddingProvider
+from shared_configs.enums import EmbedTextType
+
+
 MODEL_WARM_UP_STRING = "hi " * 512
 INFORMATION_CONTENT_MODEL_WARM_UP_STRING = "hi " * 16
+DEFAULT_OPENAI_MODEL = "text-embedding-3-small"
+DEFAULT_COHERE_MODEL = "embed-english-light-v3.0"
+DEFAULT_VOYAGE_MODEL = "voyage-large-2-instruct"
+DEFAULT_VERTEX_MODEL = "text-embedding-005"
+
+
+class EmbeddingModelTextType:
+    PROVIDER_TEXT_TYPE_MAP = {
+        EmbeddingProvider.COHERE: {
+            EmbedTextType.QUERY: "search_query",
+            EmbedTextType.PASSAGE: "search_document",
+        },
+        EmbeddingProvider.VOYAGE: {
+            EmbedTextType.QUERY: "query",
+            EmbedTextType.PASSAGE: "document",
+        },
+        EmbeddingProvider.GOOGLE: {
+            EmbedTextType.QUERY: "RETRIEVAL_QUERY",
+            EmbedTextType.PASSAGE: "RETRIEVAL_DOCUMENT",
+        },
+    }
+
+    @staticmethod
+    def get_type(provider: EmbeddingProvider, text_type: EmbedTextType) -> str:
+        return EmbeddingModelTextType.PROVIDER_TEXT_TYPE_MAP[provider][text_type]
 
 
 class GPUStatus:

backend/model_server/encoders.py

@@ -1,30 +1,55 @@
 import asyncio
+import json
 import time
-from typing import Any
+from types import TracebackType
+from typing import cast
 from typing import Optional
 
+import aioboto3  # type: ignore
+import httpx
+import openai
+import vertexai  # type: ignore
+import voyageai  # type: ignore
+from cohere import AsyncClient as CohereAsyncClient
 from fastapi import APIRouter
 from fastapi import HTTPException
 from fastapi import Request
+from google.oauth2 import service_account  # type: ignore
+from litellm import aembedding
 from litellm.exceptions import RateLimitError
+from retry import retry
 from sentence_transformers import CrossEncoder  # type: ignore
 from sentence_transformers import SentenceTransformer  # type: ignore
+from vertexai.language_models import TextEmbeddingInput  # type: ignore
+from vertexai.language_models import TextEmbeddingModel  # type: ignore
 
+from model_server.constants import DEFAULT_COHERE_MODEL
+from model_server.constants import DEFAULT_OPENAI_MODEL
+from model_server.constants import DEFAULT_VERTEX_MODEL
+from model_server.constants import DEFAULT_VOYAGE_MODEL
+from model_server.constants import EmbeddingModelTextType
+from model_server.constants import EmbeddingProvider
+from model_server.utils import pass_aws_key
 from model_server.utils import simple_log_function_time
 from onyx.utils.logger import setup_logger
+from shared_configs.configs import API_BASED_EMBEDDING_TIMEOUT
 from shared_configs.configs import INDEXING_ONLY
+from shared_configs.configs import OPENAI_EMBEDDING_TIMEOUT
+from shared_configs.configs import VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE
 from shared_configs.enums import EmbedTextType
+from shared_configs.enums import RerankerProvider
 from shared_configs.model_server_models import Embedding
 from shared_configs.model_server_models import EmbedRequest
 from shared_configs.model_server_models import EmbedResponse
 from shared_configs.model_server_models import RerankRequest
 from shared_configs.model_server_models import RerankResponse
+from shared_configs.utils import batch_list
+
 
 logger = setup_logger()
 
 router = APIRouter(prefix="/encoder")
 
-
 _GLOBAL_MODELS_DICT: dict[str, "SentenceTransformer"] = {}
 _RERANK_MODEL: Optional["CrossEncoder"] = None
 
@@ -32,6 +57,315 @@ _RERANK_MODEL: Optional["CrossEncoder"] = None
 _RETRY_DELAY = 10 if INDEXING_ONLY else 0.1
 _RETRY_TRIES = 10 if INDEXING_ONLY else 2
 
+# OpenAI only allows 2048 embeddings to be computed at once
+_OPENAI_MAX_INPUT_LEN = 2048
+# Cohere allows up to 96 embeddings in a single embedding calling
+_COHERE_MAX_INPUT_LEN = 96
+
+# Authentication error string constants
+_AUTH_ERROR_401 = "401"
+_AUTH_ERROR_UNAUTHORIZED = "unauthorized"
+_AUTH_ERROR_INVALID_API_KEY = "invalid api key"
+_AUTH_ERROR_PERMISSION = "permission"
+
+
+def is_authentication_error(error: Exception) -> bool:
+    """Check if an exception is related to authentication issues.
+
+    Args:
+        error: The exception to check
+
+    Returns:
+        bool: True if the error appears to be authentication-related
+    """
+    error_str = str(error).lower()
+    return (
+        _AUTH_ERROR_401 in error_str
+        or _AUTH_ERROR_UNAUTHORIZED in error_str
+        or _AUTH_ERROR_INVALID_API_KEY in error_str
+        or _AUTH_ERROR_PERMISSION in error_str
+    )
+
+
+def format_embedding_error(
+    error: Exception,
+    service_name: str,
+    model: str | None,
+    provider: EmbeddingProvider,
+    sanitized_api_key: str | None = None,
+    status_code: int | None = None,
+) -> str:
+    """
+    Format a standardized error string for embedding errors.
+    """
+    detail = f"Status {status_code}" if status_code else f"{type(error)}"
+
+    return (
+        f"{'HTTP error' if status_code else 'Exception'} embedding text with {service_name} - {detail}: "
+        f"Model: {model} "
+        f"Provider: {provider} "
+        f"API Key: {sanitized_api_key} "
+        f"Exception: {error}"
+    )
+
+
+# Custom exception for authentication errors
+class AuthenticationError(Exception):
+    """Raised when authentication fails with a provider."""
+
+    def __init__(self, provider: str, message: str = "API key is invalid or expired"):
+        self.provider = provider
+        self.message = message
+        super().__init__(f"{provider} authentication failed: {message}")
+
+
+class CloudEmbedding:
+    def __init__(
+        self,
+        api_key: str,
+        provider: EmbeddingProvider,
+        api_url: str | None = None,
+        api_version: str | None = None,
+        timeout: int = API_BASED_EMBEDDING_TIMEOUT,
+    ) -> None:
+        self.provider = provider
+        self.api_key = api_key
+        self.api_url = api_url
+        self.api_version = api_version
+        self.timeout = timeout
+        self.http_client = httpx.AsyncClient(timeout=timeout)
+        self._closed = False
+        self.sanitized_api_key = api_key[:4] + "********" + api_key[-4:]
+
+    async def _embed_openai(
+        self, texts: list[str], model: str | None, reduced_dimension: int | None
+    ) -> list[Embedding]:
+        if not model:
+            model = DEFAULT_OPENAI_MODEL
+
+        # Use the OpenAI specific timeout for this one
+        client = openai.AsyncOpenAI(
+            api_key=self.api_key, timeout=OPENAI_EMBEDDING_TIMEOUT
+        )
+
+        final_embeddings: list[Embedding] = []
+
+        for text_batch in batch_list(texts, _OPENAI_MAX_INPUT_LEN):
+            response = await client.embeddings.create(
+                input=text_batch,
+                model=model,
+                dimensions=reduced_dimension or openai.NOT_GIVEN,
+            )
+            final_embeddings.extend(
+                [embedding.embedding for embedding in response.data]
+            )
+        return final_embeddings
+
+    async def _embed_cohere(
+        self, texts: list[str], model: str | None, embedding_type: str
+    ) -> list[Embedding]:
+        if not model:
+            model = DEFAULT_COHERE_MODEL
+
+        client = CohereAsyncClient(api_key=self.api_key)
+
+        final_embeddings: list[Embedding] = []
+        for text_batch in batch_list(texts, _COHERE_MAX_INPUT_LEN):
+            # Does not use the same tokenizer as the Onyx API server but it's approximately the same
+            # empirically it's only off by a very few tokens so it's not a big deal
+            response = await client.embed(
+                texts=text_batch,
+                model=model,
+                input_type=embedding_type,
+                truncate="END",
+            )
+            final_embeddings.extend(cast(list[Embedding], response.embeddings))
+        return final_embeddings
+
+    async def _embed_voyage(
+        self, texts: list[str], model: str | None, embedding_type: str
+    ) -> list[Embedding]:
+        if not model:
+            model = DEFAULT_VOYAGE_MODEL
+
+        client = voyageai.AsyncClient(
+            api_key=self.api_key, timeout=API_BASED_EMBEDDING_TIMEOUT
+        )
+
+        response = await client.embed(
+            texts=texts,
+            model=model,
+            input_type=embedding_type,
+            truncation=True,
+        )
+        return response.embeddings
+
+    async def _embed_azure(
+        self, texts: list[str], model: str | None
+    ) -> list[Embedding]:
+        response = await aembedding(
+            model=model,
+            input=texts,
+            timeout=API_BASED_EMBEDDING_TIMEOUT,
+            api_key=self.api_key,
+            api_base=self.api_url,
+            api_version=self.api_version,
+        )
+        embeddings = [embedding["embedding"] for embedding in response.data]
+        return embeddings
+
+    async def _embed_vertex(
+        self, texts: list[str], model: str | None, embedding_type: str
+    ) -> list[Embedding]:
+        if not model:
+            model = DEFAULT_VERTEX_MODEL
+
+        credentials = service_account.Credentials.from_service_account_info(
+            json.loads(self.api_key)
+        )
+        project_id = json.loads(self.api_key)["project_id"]
+        vertexai.init(project=project_id, credentials=credentials)
+        client = TextEmbeddingModel.from_pretrained(model)
+
+        inputs = [TextEmbeddingInput(text, embedding_type) for text in texts]
+
+        # Split into batches of 25 texts
+        max_texts_per_batch = VERTEXAI_EMBEDDING_LOCAL_BATCH_SIZE
+        batches = [
+            inputs[i : i + max_texts_per_batch]
+            for i in range(0, len(inputs), max_texts_per_batch)
+        ]
+
+        # Dispatch all embedding calls asynchronously at once
+        tasks = [
+            client.get_embeddings_async(batch, auto_truncate=True) for batch in batches
+        ]
+
+        # Wait for all tasks to complete in parallel
+        results = await asyncio.gather(*tasks)
+
+        return [embedding.values for batch in results for embedding in batch]
+
+    async def _embed_litellm_proxy(
+        self, texts: list[str], model_name: str | None
+    ) -> list[Embedding]:
+        if not model_name:
+            raise ValueError("Model name is required for LiteLLM proxy embedding.")
+
+        if not self.api_url:
+            raise ValueError("API URL is required for LiteLLM proxy embedding.")
+
+        headers = (
+            {} if not self.api_key else {"Authorization": f"Bearer {self.api_key}"}
+        )
+
+        response = await self.http_client.post(
+            self.api_url,
+            json={
+                "model": model_name,
+                "input": texts,
+            },
+            headers=headers,
+        )
+        response.raise_for_status()
+        result = response.json()
+        return [embedding["embedding"] for embedding in result["data"]]
+
+    @retry(tries=_RETRY_TRIES, delay=_RETRY_DELAY)
+    async def embed(
+        self,
+        *,
+        texts: list[str],
+        text_type: EmbedTextType,
+        model_name: str | None = None,
+        deployment_name: str | None = None,
+        reduced_dimension: int | None = None,
+    ) -> list[Embedding]:
+        try:
+            if self.provider == EmbeddingProvider.OPENAI:
+                return await self._embed_openai(texts, model_name, reduced_dimension)
+            elif self.provider == EmbeddingProvider.AZURE:
+                return await self._embed_azure(texts, f"azure/{deployment_name}")
+            elif self.provider == EmbeddingProvider.LITELLM:
+                return await self._embed_litellm_proxy(texts, model_name)
+
+            embedding_type = EmbeddingModelTextType.get_type(self.provider, text_type)
+            if self.provider == EmbeddingProvider.COHERE:
+                return await self._embed_cohere(texts, model_name, embedding_type)
+            elif self.provider == EmbeddingProvider.VOYAGE:
+                return await self._embed_voyage(texts, model_name, embedding_type)
+            elif self.provider == EmbeddingProvider.GOOGLE:
+                return await self._embed_vertex(texts, model_name, embedding_type)
+            else:
+                raise ValueError(f"Unsupported provider: {self.provider}")
+        except openai.AuthenticationError:
+            raise AuthenticationError(provider="OpenAI")
+        except httpx.HTTPStatusError as e:
+            if e.response.status_code == 401:
+                raise AuthenticationError(provider=str(self.provider))
+
+            error_string = format_embedding_error(
+                e,
+                str(self.provider),
+                model_name or deployment_name,
+                self.provider,
+                sanitized_api_key=self.sanitized_api_key,
+                status_code=e.response.status_code,
+            )
+            logger.error(error_string)
+            logger.debug(f"Exception texts: {texts}")
+
+            raise RuntimeError(error_string)
+        except Exception as e:
+            if is_authentication_error(e):
+                raise AuthenticationError(provider=str(self.provider))
+
+            error_string = format_embedding_error(
+                e,
+                str(self.provider),
+                model_name or deployment_name,
+                self.provider,
+                sanitized_api_key=self.sanitized_api_key,
+            )
+            logger.error(error_string)
+            logger.debug(f"Exception texts: {texts}")
+
+            raise RuntimeError(error_string)
+
+    @staticmethod
+    def create(
+        api_key: str,
+        provider: EmbeddingProvider,
+        api_url: str | None = None,
+        api_version: str | None = None,
+    ) -> "CloudEmbedding":
+        logger.debug(f"Creating Embedding instance for provider: {provider}")
+        return CloudEmbedding(api_key, provider, api_url, api_version)
+
+    async def aclose(self) -> None:
+        """Explicitly close the client."""
+        if not self._closed:
+            await self.http_client.aclose()
+            self._closed = True
+
+    async def __aenter__(self) -> "CloudEmbedding":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+    def __del__(self) -> None:
+        """Finalizer to warn about unclosed clients."""
+        if not self._closed:
+            logger.warning(
+                "CloudEmbedding was not properly closed. Use 'async with' or call aclose()"
+            )
+
 
 def get_embedding_model(
     model_name: str,
@@ -70,34 +404,20 @@ def get_local_reranking_model(
     return _RERANK_MODEL
 
 
-ENCODING_RETRIES = 3
-ENCODING_RETRY_DELAY = 0.1
-
-
-def _concurrent_embedding(
-    texts: list[str], model: "SentenceTransformer", normalize_embeddings: bool
-) -> Any:
-    """Synchronous wrapper for concurrent_embedding to use with run_in_executor."""
-    for _ in range(ENCODING_RETRIES):
-        try:
-            return model.encode(texts, normalize_embeddings=normalize_embeddings)
-        except RuntimeError as e:
-            # There is a concurrency bug in the SentenceTransformer library that causes
-            # the model to fail to encode texts. It's pretty rare and we want to allow
-            # concurrent embedding, hence we retry (the specific error is
-            # "RuntimeError: Already borrowed" and occurs in the transformers library)
-            logger.error(f"Error encoding texts, retrying: {e}")
-            time.sleep(ENCODING_RETRY_DELAY)
-    return model.encode(texts, normalize_embeddings=normalize_embeddings)
-
-
 @simple_log_function_time()
 async def embed_text(
     texts: list[str],
+    text_type: EmbedTextType,
     model_name: str | None,
+    deployment_name: str | None,
     max_context_length: int,
     normalize_embeddings: bool,
+    api_key: str | None,
+    provider_type: EmbeddingProvider | None,
     prefix: str | None,
+    api_url: str | None,
+    api_version: str | None,
+    reduced_dimension: int | None,
     gpu_type: str = "UNKNOWN",
 ) -> list[Embedding]:
     if not all(texts):
@@ -114,10 +434,52 @@ async def embed_text(
     for text in texts:
         total_chars += len(text)
 
-    # Only local models should call this function now
-    # API providers should go directly to API server
+    if provider_type is not None:
+        logger.info(
+            f"Embedding {len(texts)} texts with {total_chars} total characters with provider: {provider_type}"
+        )
 
-    if model_name is not None:
+        if api_key is None:
+            logger.error("API key not provided for cloud model")
+            raise RuntimeError("API key not provided for cloud model")
+
+        if prefix:
+            logger.warning("Prefix provided for cloud model, which is not supported")
+            raise ValueError(
+                "Prefix string is not valid for cloud models. "
+                "Cloud models take an explicit text type instead."
+            )
+
+        async with CloudEmbedding(
+            api_key=api_key,
+            provider=provider_type,
+            api_url=api_url,
+            api_version=api_version,
+        ) as cloud_model:
+            embeddings = await cloud_model.embed(
+                texts=texts,
+                model_name=model_name,
+                deployment_name=deployment_name,
+                text_type=text_type,
+                reduced_dimension=reduced_dimension,
+            )
+
+        if any(embedding is None for embedding in embeddings):
+            error_message = "Embeddings contain None values\n"
+            error_message += "Corresponding texts:\n"
+            error_message += "\n".join(texts)
+            logger.error(error_message)
+            raise ValueError(error_message)
+
+        elapsed = time.monotonic() - start
+        logger.info(
+            f"event=embedding_provider "
+            f"texts={len(texts)} "
+            f"chars={total_chars} "
+            f"provider={provider_type} "
+            f"elapsed={elapsed:.2f}"
+        )
+    elif model_name is not None:
         logger.info(
             f"Embedding {len(texts)} texts with {total_chars} total characters with local model: {model_name}"
         )
@@ -130,8 +492,8 @@ async def embed_text(
         # Run CPU-bound embedding in a thread pool
         embeddings_vectors = await asyncio.get_event_loop().run_in_executor(
             None,
-            lambda: _concurrent_embedding(
-                prefixed_texts, local_model, normalize_embeddings
+            lambda: local_model.encode(
+                prefixed_texts, normalize_embeddings=normalize_embeddings
             ),
         )
         embeddings = [
@@ -153,8 +515,10 @@ async def embed_text(
             f"elapsed={elapsed:.2f}"
         )
     else:
-        logger.error("Model name not specified for embedding")
-        raise ValueError("Model name must be provided to run embeddings.")
+        logger.error("Neither model name nor provider specified for embedding")
+        raise ValueError(
+            "Either model name or provider must be provided to run embeddings."
+        )
 
     return embeddings
 
@@ -169,6 +533,77 @@ async def local_rerank(query: str, docs: list[str], model_name: str) -> list[flo
     )
 
 
+async def cohere_rerank_api(
+    query: str, docs: list[str], model_name: str, api_key: str
+) -> list[float]:
+    cohere_client = CohereAsyncClient(api_key=api_key)
+    response = await cohere_client.rerank(query=query, documents=docs, model=model_name)
+    results = response.results
+    sorted_results = sorted(results, key=lambda item: item.index)
+    return [result.relevance_score for result in sorted_results]
+
+
+async def cohere_rerank_aws(
+    query: str,
+    docs: list[str],
+    model_name: str,
+    region_name: str,
+    aws_access_key_id: str,
+    aws_secret_access_key: str,
+) -> list[float]:
+    session = aioboto3.Session(
+        aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key
+    )
+    async with session.client(
+        "bedrock-runtime", region_name=region_name
+    ) as bedrock_client:
+        body = json.dumps(
+            {
+                "query": query,
+                "documents": docs,
+                "api_version": 2,
+            }
+        )
+        # Invoke the Bedrock model asynchronously
+        response = await bedrock_client.invoke_model(
+            modelId=model_name,
+            accept="application/json",
+            contentType="application/json",
+            body=body,
+        )
+
+        # Read the response asynchronously
+        response_body = json.loads(await response["body"].read())
+
+        # Extract and sort the results
+        results = response_body.get("results", [])
+        sorted_results = sorted(results, key=lambda item: item["index"])
+
+        return [result["relevance_score"] for result in sorted_results]
+
+
+async def litellm_rerank(
+    query: str, docs: list[str], api_url: str, model_name: str, api_key: str | None
+) -> list[float]:
+    headers = {} if not api_key else {"Authorization": f"Bearer {api_key}"}
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            api_url,
+            json={
+                "model": model_name,
+                "query": query,
+                "documents": docs,
+            },
+            headers=headers,
+        )
+        response.raise_for_status()
+        result = response.json()
+        return [
+            item["relevance_score"]
+            for item in sorted(result["results"], key=lambda x: x["index"])
+        ]
+
+
 @router.post("/bi-encoder-embed")
 async def route_bi_encoder_embed(
     request: Request,
@@ -180,13 +615,6 @@ async def route_bi_encoder_embed(
 async def process_embed_request(
     embed_request: EmbedRequest, gpu_type: str = "UNKNOWN"
 ) -> EmbedResponse:
-    # Only local models should use this endpoint - API providers should make direct API calls
-    if embed_request.provider_type is not None:
-        raise ValueError(
-            f"Model server embedding endpoint should only be used for local models. "
-            f"API provider '{embed_request.provider_type}' should make direct API calls instead."
-        )
-
     if not embed_request.texts:
         raise HTTPException(status_code=400, detail="No texts to be embedded")
 
@@ -204,12 +632,26 @@ async def process_embed_request(
         embeddings = await embed_text(
             texts=embed_request.texts,
             model_name=embed_request.model_name,
+            deployment_name=embed_request.deployment_name,
             max_context_length=embed_request.max_context_length,
             normalize_embeddings=embed_request.normalize_embeddings,
+            api_key=embed_request.api_key,
+            provider_type=embed_request.provider_type,
+            text_type=embed_request.text_type,
+            api_url=embed_request.api_url,
+            api_version=embed_request.api_version,
+            reduced_dimension=embed_request.reduced_dimension,
             prefix=prefix,
             gpu_type=gpu_type,
         )
         return EmbedResponse(embeddings=embeddings)
+    except AuthenticationError as e:
+        # Handle authentication errors consistently
+        logger.error(f"Authentication error: {e.provider}")
+        raise HTTPException(
+            status_code=401,
+            detail=f"Authentication failed: {e.message}",
+        )
     except RateLimitError as e:
         raise HTTPException(
             status_code=429,
@@ -227,13 +669,6 @@ async def process_embed_request(
 @router.post("/cross-encoder-scores")
 async def process_rerank_request(rerank_request: RerankRequest) -> RerankResponse:
     """Cross encoders can be purely black box from the app perspective"""
-    # Only local models should use this endpoint - API providers should make direct API calls
-    if rerank_request.provider_type is not None:
-        raise ValueError(
-            f"Model server reranking endpoint should only be used for local models. "
-            f"API provider '{rerank_request.provider_type}' should make direct API calls instead."
-        )
-
     if INDEXING_ONLY:
         raise RuntimeError("Indexing model server should not call intent endpoint")
 
@@ -245,13 +680,55 @@ async def process_rerank_request(rerank_request: RerankRequest) -> RerankRespons
         raise ValueError("Empty documents cannot be reranked.")
 
     try:
-        # At this point, provider_type is None, so handle local reranking
-        sim_scores = await local_rerank(
-            query=rerank_request.query,
-            docs=rerank_request.documents,
-            model_name=rerank_request.model_name,
-        )
-        return RerankResponse(scores=sim_scores)
+        if rerank_request.provider_type is None:
+            sim_scores = await local_rerank(
+                query=rerank_request.query,
+                docs=rerank_request.documents,
+                model_name=rerank_request.model_name,
+            )
+            return RerankResponse(scores=sim_scores)
+        elif rerank_request.provider_type == RerankerProvider.LITELLM:
+            if rerank_request.api_url is None:
+                raise ValueError("API URL is required for LiteLLM reranking.")
+
+            sim_scores = await litellm_rerank(
+                query=rerank_request.query,
+                docs=rerank_request.documents,
+                api_url=rerank_request.api_url,
+                model_name=rerank_request.model_name,
+                api_key=rerank_request.api_key,
+            )
+
+            return RerankResponse(scores=sim_scores)
+
+        elif rerank_request.provider_type == RerankerProvider.COHERE:
+            if rerank_request.api_key is None:
+                raise RuntimeError("Cohere Rerank Requires an API Key")
+            sim_scores = await cohere_rerank_api(
+                query=rerank_request.query,
+                docs=rerank_request.documents,
+                model_name=rerank_request.model_name,
+                api_key=rerank_request.api_key,
+            )
+            return RerankResponse(scores=sim_scores)
+
+        elif rerank_request.provider_type == RerankerProvider.BEDROCK:
+            if rerank_request.api_key is None:
+                raise RuntimeError("Bedrock Rerank Requires an API Key")
+            aws_access_key_id, aws_secret_access_key, aws_region = pass_aws_key(
+                rerank_request.api_key
+            )
+            sim_scores = await cohere_rerank_aws(
+                query=rerank_request.query,
+                docs=rerank_request.documents,
+                model_name=rerank_request.model_name,
+                region_name=aws_region,
+                aws_access_key_id=aws_access_key_id,
+                aws_secret_access_key=aws_secret_access_key,
+            )
+            return RerankResponse(scores=sim_scores)
+        else:
+            raise ValueError(f"Unsupported provider: {rerank_request.provider_type}")
 
     except Exception as e:
         logger.exception(f"Error during reranking process:\n{str(e)}")

backend/model_server/main.py

@@ -34,8 +34,8 @@ from shared_configs.configs import SENTRY_DSN
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
 
-HF_CACHE_PATH = Path(".cache/huggingface")
-TEMP_HF_CACHE_PATH = Path(".cache/temp_huggingface")
+HF_CACHE_PATH = Path(os.path.expanduser("~")) / ".cache/huggingface"
+TEMP_HF_CACHE_PATH = Path(os.path.expanduser("~")) / ".cache/temp_huggingface"
 
 transformer_logging.set_verbosity_error()
 

backend/model_server/utils.py

@@ -70,3 +70,32 @@ def get_gpu_type() -> str:
         return GPUStatus.MAC_MPS
 
     return GPUStatus.NONE
+
+
+def pass_aws_key(api_key: str) -> tuple[str, str, str]:
+    """Parse AWS API key string into components.
+
+    Args:
+        api_key: String in format 'aws_ACCESSKEY_SECRETKEY_REGION'
+
+    Returns:
+        Tuple of (access_key, secret_key, region)
+
+    Raises:
+        ValueError: If key format is invalid
+    """
+    if not api_key.startswith("aws"):
+        raise ValueError("API key must start with 'aws' prefix")
+
+    parts = api_key.split("_")
+    if len(parts) != 4:
+        raise ValueError(
+            f"API key must be in format 'aws_ACCESSKEY_SECRETKEY_REGION', got {len(parts) - 1} parts"
+            "this is an onyx specific format for formatting the aws secrets for bedrock"
+        )
+
+    try:
+        _, aws_access_key_id, aws_secret_access_key, aws_region = parts
+        return aws_access_key_id, aws_secret_access_key, aws_region
+    except Exception as e:
+        raise ValueError(f"Failed to parse AWS key components: {str(e)}")

backend/onyx/agents/agent_search/basic/graph_builder.py

@@ -0,0 +1,97 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agents.agent_search.basic.states import BasicInput
+from onyx.agents.agent_search.basic.states import BasicOutput
+from onyx.agents.agent_search.basic.states import BasicState
+from onyx.agents.agent_search.orchestration.nodes.call_tool import call_tool
+from onyx.agents.agent_search.orchestration.nodes.choose_tool import choose_tool
+from onyx.agents.agent_search.orchestration.nodes.prepare_tool_input import (
+    prepare_tool_input,
+)
+from onyx.agents.agent_search.orchestration.nodes.use_tool_response import (
+    basic_use_tool_response,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def basic_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=BasicState,
+        input=BasicInput,
+        output=BasicOutput,
+    )
+
+    ### Add nodes ###
+
+    graph.add_node(
+        node="prepare_tool_input",
+        action=prepare_tool_input,
+    )
+
+    graph.add_node(
+        node="choose_tool",
+        action=choose_tool,
+    )
+
+    graph.add_node(
+        node="call_tool",
+        action=call_tool,
+    )
+
+    graph.add_node(
+        node="basic_use_tool_response",
+        action=basic_use_tool_response,
+    )
+
+    ### Add edges ###
+
+    graph.add_edge(start_key=START, end_key="prepare_tool_input")
+
+    graph.add_edge(start_key="prepare_tool_input", end_key="choose_tool")
+
+    graph.add_conditional_edges("choose_tool", should_continue, ["call_tool", END])
+
+    graph.add_edge(
+        start_key="call_tool",
+        end_key="basic_use_tool_response",
+    )
+
+    graph.add_edge(
+        start_key="basic_use_tool_response",
+        end_key=END,
+    )
+
+    return graph
+
+
+def should_continue(state: BasicState) -> str:
+    return (
+        # If there are no tool calls, basic graph already streamed the answer
+        END
+        if state.tool_choice is None
+        else "call_tool"
+    )
+
+
+if __name__ == "__main__":
+    from onyx.db.engine.sql_engine import get_session_with_current_tenant
+    from onyx.context.search.models import SearchRequest
+    from onyx.llm.factory import get_default_llms
+    from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config
+
+    graph = basic_graph_builder()
+    compiled_graph = graph.compile()
+    input = BasicInput(unused=True)
+    primary_llm, fast_llm = get_default_llms()
+    with get_session_with_current_tenant() as db_session:
+        config, _ = get_test_config(
+            db_session=db_session,
+            primary_llm=primary_llm,
+            fast_llm=fast_llm,
+            search_request=SearchRequest(query="How does onyx use FastAPI?"),
+        )
+        compiled_graph.invoke(input, config={"metadata": {"config": config}})

backend/onyx/agents/agent_search/basic/states.py

@@ -0,0 +1,35 @@
+from typing import TypedDict
+
+from langchain_core.messages import AIMessageChunk
+from pydantic import BaseModel
+
+from onyx.agents.agent_search.orchestration.states import ToolCallUpdate
+from onyx.agents.agent_search.orchestration.states import ToolChoiceInput
+from onyx.agents.agent_search.orchestration.states import ToolChoiceUpdate
+
+# States contain values that change over the course of graph execution,
+# Config is for values that are set at the start and never change.
+# If you are using a value from the config and realize it needs to change,
+# you should add it to the state and use/update the version in the state.
+
+
+## Graph Input State
+class BasicInput(BaseModel):
+    # Langgraph needs a nonempty input, but we pass in all static
+    # data through a RunnableConfig.
+    unused: bool = True
+
+
+## Graph Output State
+class BasicOutput(TypedDict):
+    tool_call_chunk: AIMessageChunk
+
+
+## Graph State
+class BasicState(
+    BasicInput,
+    ToolChoiceInput,
+    ToolCallUpdate,
+    ToolChoiceUpdate,
+):
+    pass

backend/onyx/agents/agent_search/basic/utils.py

@@ -0,0 +1,64 @@
+from collections.abc import Iterator
+from typing import cast
+
+from langchain_core.messages import AIMessageChunk
+from langchain_core.messages import BaseMessage
+from langgraph.types import StreamWriter
+
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import LlmDoc
+from onyx.chat.stream_processing.answer_response_handler import AnswerResponseHandler
+from onyx.chat.stream_processing.answer_response_handler import CitationResponseHandler
+from onyx.chat.stream_processing.answer_response_handler import (
+    PassThroughAnswerResponseHandler,
+)
+from onyx.chat.stream_processing.utils import map_document_id_order
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def process_llm_stream(
+    messages: Iterator[BaseMessage],
+    should_stream_answer: bool,
+    writer: StreamWriter,
+    final_search_results: list[LlmDoc] | None = None,
+    displayed_search_results: list[LlmDoc] | None = None,
+) -> AIMessageChunk:
+    tool_call_chunk = AIMessageChunk(content="")
+
+    if final_search_results and displayed_search_results:
+        answer_handler: AnswerResponseHandler = CitationResponseHandler(
+            context_docs=final_search_results,
+            final_doc_id_to_rank_map=map_document_id_order(final_search_results),
+            display_doc_id_to_rank_map=map_document_id_order(displayed_search_results),
+        )
+    else:
+        answer_handler = PassThroughAnswerResponseHandler()
+
+    full_answer = ""
+    # This stream will be the llm answer if no tool is chosen. When a tool is chosen,
+    # the stream will contain AIMessageChunks with tool call information.
+    for message in messages:
+
+        answer_piece = message.content
+        if not isinstance(answer_piece, str):
+            # this is only used for logging, so fine to
+            # just add the string representation
+            answer_piece = str(answer_piece)
+        full_answer += answer_piece
+
+        if isinstance(message, AIMessageChunk) and (
+            message.tool_call_chunks or message.tool_calls
+        ):
+            tool_call_chunk += message  # type: ignore
+        elif should_stream_answer:
+            for response_part in answer_handler.handle_response_part(message, []):
+                write_custom_event(
+                    "basic_response",
+                    response_part,
+                    writer,
+                )
+
+    logger.debug(f"Full answer: {full_answer}")
+    return cast(AIMessageChunk, tool_call_chunk)

backend/onyx/agents/agent_search/core_state.py

@@ -10,7 +10,6 @@ class CoreState(BaseModel):
     """
 
     log_messages: Annotated[list[str], add] = []
-    current_step_nr: int = 1
 
 
 class SubgraphCoreState(BaseModel):

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a1_search_objects.py

@@ -14,6 +14,8 @@ from onyx.agents.agent_search.models import GraphConfig
 from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     trim_prompt_piece,
 )
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import AgentAnswerPiece
 from onyx.configs.constants import DocumentSource
 from onyx.prompts.agents.dc_prompts import DC_OBJECT_NO_BASE_DATA_EXTRACTION_PROMPT
 from onyx.prompts.agents.dc_prompts import DC_OBJECT_SEPARATOR
@@ -137,6 +139,17 @@ def search_objects(
     except Exception as e:
         raise ValueError(f"Error in search_objects: {e}")
 
+    write_custom_event(
+        "initial_agent_answer",
+        AgentAnswerPiece(
+            answer_piece=" Researching the individual objects for each source type... ",
+            level=0,
+            level_question_num=0,
+            answer_type="agent_level_answer",
+        ),
+        writer,
+    )
+
     return SearchSourcesObjectsUpdate(
         analysis_objects=object_list,
         analysis_sources=document_sources,

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a3_structure_research_by_object.py

@@ -9,6 +9,8 @@ from onyx.agents.agent_search.dc_search_analysis.states import MainState
 from onyx.agents.agent_search.dc_search_analysis.states import (
     ObjectResearchInformationUpdate,
 )
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import AgentAnswerPiece
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -21,6 +23,17 @@ def structure_research_by_object(
     LangGraph node to start the agentic search process.
     """
 
+    write_custom_event(
+        "initial_agent_answer",
+        AgentAnswerPiece(
+            answer_piece=" consolidating the information across source types for each object...",
+            level=0,
+            level_question_num=0,
+            answer_type="agent_level_answer",
+        ),
+        writer,
+    )
+
     object_source_research_results = state.object_source_research_results
 
     object_research_information_results: List[Dict[str, str]] = []

backend/onyx/agents/agent_search/dc_search_analysis/nodes/a5_consolidate_research.py

@@ -12,6 +12,8 @@ from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
     trim_prompt_piece,
 )
 from onyx.agents.agent_search.shared_graph_utils.llm import stream_llm_answer
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import AgentAnswerPiece
 from onyx.prompts.agents.dc_prompts import DC_FORMATTING_NO_BASE_DATA_PROMPT
 from onyx.prompts.agents.dc_prompts import DC_FORMATTING_WITH_BASE_DATA_PROMPT
 from onyx.utils.logger import setup_logger
@@ -31,6 +33,17 @@ def consolidate_research(
 
     search_tool = graph_config.tooling.search_tool
 
+    write_custom_event(
+        "initial_agent_answer",
+        AgentAnswerPiece(
+            answer_piece=" generating the answer\n\n\n",
+            level=0,
+            level_question_num=0,
+            answer_type="agent_level_answer",
+        ),
+        writer,
+    )
+
     if search_tool is None or graph_config.inputs.persona is None:
         raise ValueError("Search tool and persona must be provided for DivCon search")
 

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/edges.py

@@ -0,0 +1,31 @@
+from collections.abc import Hashable
+from datetime import datetime
+
+from langgraph.types import Send
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionAnsweringInput,
+)
+from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def send_to_expanded_retrieval(state: SubQuestionAnsweringInput) -> Send | Hashable:
+    """
+    LangGraph edge to send a sub-question to the expanded retrieval.
+    """
+    edge_start_time = datetime.now()
+
+    return Send(
+        "initial_sub_question_expanded_retrieval",
+        ExpandedRetrievalInput(
+            question=state.question,
+            base_search=False,
+            sub_question_id=state.question_id,
+            log_messages=[f"{edge_start_time} -- Sending to expanded retrieval"],
+        ),
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/graph_builder.py

@@ -0,0 +1,137 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.edges import (
+    send_to_expanded_retrieval,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.nodes.check_sub_answer import (
+    check_sub_answer,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.nodes.format_sub_answer import (
+    format_sub_answer,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.nodes.generate_sub_answer import (
+    generate_sub_answer,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.nodes.ingest_retrieved_documents import (
+    ingest_retrieved_documents,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionState,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionAnsweringInput,
+)
+from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_query_graph_builder() -> StateGraph:
+    """
+    LangGraph sub-graph builder for the initial individual sub-answer generation.
+    """
+    graph = StateGraph(
+        state_schema=AnswerQuestionState,
+        input=SubQuestionAnsweringInput,
+        output=AnswerQuestionOutput,
+    )
+
+    ### Add nodes ###
+
+    # The sub-graph that executes the expanded retrieval process for a sub-question
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="initial_sub_question_expanded_retrieval",
+        action=expanded_retrieval,
+    )
+
+    # The node that ingests the retrieved documents and puts them into the proper
+    # state keys.
+    graph.add_node(
+        node="ingest_retrieval",
+        action=ingest_retrieved_documents,
+    )
+
+    # The node that generates the sub-answer
+    graph.add_node(
+        node="generate_sub_answer",
+        action=generate_sub_answer,
+    )
+
+    # The node that checks the sub-answer
+    graph.add_node(
+        node="answer_check",
+        action=check_sub_answer,
+    )
+
+    # The node that formats the sub-answer for the following initial answer generation
+    graph.add_node(
+        node="format_answer",
+        action=format_sub_answer,
+    )
+
+    ### Add edges ###
+
+    graph.add_conditional_edges(
+        source=START,
+        path=send_to_expanded_retrieval,
+        path_map=["initial_sub_question_expanded_retrieval"],
+    )
+    graph.add_edge(
+        start_key="initial_sub_question_expanded_retrieval",
+        end_key="ingest_retrieval",
+    )
+    graph.add_edge(
+        start_key="ingest_retrieval",
+        end_key="generate_sub_answer",
+    )
+    graph.add_edge(
+        start_key="generate_sub_answer",
+        end_key="answer_check",
+    )
+    graph.add_edge(
+        start_key="answer_check",
+        end_key="format_answer",
+    )
+    graph.add_edge(
+        start_key="format_answer",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine.sql_engine import get_session_with_current_tenant
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = answer_query_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+    with get_session_with_current_tenant() as db_session:
+        graph_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+        inputs = SubQuestionAnsweringInput(
+            question="what can you do with onyx?",
+            question_id="0_0",
+            log_messages=[],
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            config={"configurable": {"config": graph_config}},
+        ):
+            logger.debug(thing)

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/nodes/check_sub_answer.py

@@ -0,0 +1,136 @@
+from datetime import datetime
+from typing import cast
+
+from langchain_core.messages import BaseMessage
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables.config import RunnableConfig
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionState,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionAnswerCheckUpdate,
+)
+from onyx.agents.agent_search.models import GraphConfig
+from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
+    binary_string_test,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_LLM_RATELIMIT_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_LLM_TIMEOUT_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_POSITIVE_VALUE_STR,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import AgentLLMErrorType
+from onyx.agents.agent_search.shared_graph_utils.models import AgentErrorLog
+from onyx.agents.agent_search.shared_graph_utils.models import LLMNodeErrorStrings
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_langgraph_node_log_string,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import parse_question_id
+from onyx.configs.agent_configs import AGENT_MAX_TOKENS_VALIDATION
+from onyx.configs.agent_configs import AGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_CHECK
+from onyx.configs.agent_configs import AGENT_TIMEOUT_LLM_SUBANSWER_CHECK
+from onyx.llm.chat_llm import LLMRateLimitError
+from onyx.llm.chat_llm import LLMTimeoutError
+from onyx.prompts.agent_search import SUB_ANSWER_CHECK_PROMPT
+from onyx.prompts.agent_search import UNKNOWN_ANSWER
+from onyx.utils.logger import setup_logger
+from onyx.utils.threadpool_concurrency import run_with_timeout
+from onyx.utils.timing import log_function_time
+
+logger = setup_logger()
+
+_llm_node_error_strings = LLMNodeErrorStrings(
+    timeout="LLM Timeout Error. The sub-answer will be treated as 'relevant'",
+    rate_limit="LLM Rate Limit Error. The sub-answer will be treated as 'relevant'",
+    general_error="General LLM Error. The sub-answer will be treated as 'relevant'",
+)
+
+
+@log_function_time(print_only=True)
+def check_sub_answer(
+    state: AnswerQuestionState, config: RunnableConfig
+) -> SubQuestionAnswerCheckUpdate:
+    """
+    LangGraph node to check the quality of the sub-answer. The answer
+    is represented as a boolean value.
+    """
+    node_start_time = datetime.now()
+
+    level, question_num = parse_question_id(state.question_id)
+    if state.answer == UNKNOWN_ANSWER:
+        return SubQuestionAnswerCheckUpdate(
+            answer_quality=False,
+            log_messages=[
+                get_langgraph_node_log_string(
+                    graph_component="initial  - generate individual sub answer",
+                    node_name="check sub answer",
+                    node_start_time=node_start_time,
+                    result="unknown answer",
+                )
+            ],
+        )
+    msg = [
+        HumanMessage(
+            content=SUB_ANSWER_CHECK_PROMPT.format(
+                question=state.question,
+                base_answer=state.answer,
+            )
+        )
+    ]
+
+    graph_config = cast(GraphConfig, config["metadata"]["config"])
+    fast_llm = graph_config.tooling.fast_llm
+    agent_error: AgentErrorLog | None = None
+    response: BaseMessage | None = None
+    try:
+        response = run_with_timeout(
+            AGENT_TIMEOUT_LLM_SUBANSWER_CHECK,
+            fast_llm.invoke,
+            prompt=msg,
+            timeout_override=AGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_CHECK,
+            max_tokens=AGENT_MAX_TOKENS_VALIDATION,
+        )
+
+        quality_str: str = cast(str, response.content)
+        answer_quality = binary_string_test(
+            text=quality_str, positive_value=AGENT_POSITIVE_VALUE_STR
+        )
+        log_result = f"Answer quality: {quality_str}"
+
+    except (LLMTimeoutError, TimeoutError):
+        agent_error = AgentErrorLog(
+            error_type=AgentLLMErrorType.TIMEOUT,
+            error_message=AGENT_LLM_TIMEOUT_MESSAGE,
+            error_result=_llm_node_error_strings.timeout,
+        )
+        answer_quality = True
+        log_result = agent_error.error_result
+        logger.error("LLM Timeout Error - check sub answer")
+
+    except LLMRateLimitError:
+        agent_error = AgentErrorLog(
+            error_type=AgentLLMErrorType.RATE_LIMIT,
+            error_message=AGENT_LLM_RATELIMIT_MESSAGE,
+            error_result=_llm_node_error_strings.rate_limit,
+        )
+
+        answer_quality = True
+        log_result = agent_error.error_result
+        logger.error("LLM Rate Limit Error - check sub answer")
+
+    return SubQuestionAnswerCheckUpdate(
+        answer_quality=answer_quality,
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="initial  - generate individual sub answer",
+                node_name="check sub answer",
+                node_start_time=node_start_time,
+                result=log_result,
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/nodes/format_sub_answer.py

@@ -0,0 +1,30 @@
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionState,
+)
+from onyx.agents.agent_search.shared_graph_utils.models import (
+    SubQuestionAnswerResults,
+)
+
+
+def format_sub_answer(state: AnswerQuestionState) -> AnswerQuestionOutput:
+    """
+    LangGraph node to generate the sub-answer format.
+    """
+    return AnswerQuestionOutput(
+        answer_results=[
+            SubQuestionAnswerResults(
+                question=state.question,
+                question_id=state.question_id,
+                verified_high_quality=state.answer_quality,
+                answer=state.answer,
+                sub_query_retrieval_results=state.expanded_retrieval_results,
+                verified_reranked_documents=state.verified_reranked_documents,
+                context_documents=state.context_documents,
+                cited_documents=state.cited_documents,
+                sub_question_retrieval_stats=state.sub_question_retrieval_stats,
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/nodes/generate_sub_answer.py

@@ -0,0 +1,185 @@
+from datetime import datetime
+from typing import cast
+
+from langchain_core.messages import merge_message_runs
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.types import StreamWriter
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionState,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionAnswerGenerationUpdate,
+)
+from onyx.agents.agent_search.models import GraphConfig
+from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
+    build_sub_question_answer_prompt,
+)
+from onyx.agents.agent_search.shared_graph_utils.calculations import (
+    dedup_sort_inference_section_list,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_LLM_RATELIMIT_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_LLM_TIMEOUT_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AgentLLMErrorType,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    LLM_ANSWER_ERROR_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.llm import stream_llm_answer
+from onyx.agents.agent_search.shared_graph_utils.models import AgentErrorLog
+from onyx.agents.agent_search.shared_graph_utils.models import LLMNodeErrorStrings
+from onyx.agents.agent_search.shared_graph_utils.utils import get_answer_citation_ids
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_langgraph_node_log_string,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_persona_agent_prompt_expressions,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import parse_question_id
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import AgentAnswerPiece
+from onyx.chat.models import StreamStopInfo
+from onyx.chat.models import StreamStopReason
+from onyx.chat.models import StreamType
+from onyx.configs.agent_configs import AGENT_MAX_ANSWER_CONTEXT_DOCS
+from onyx.configs.agent_configs import AGENT_MAX_TOKENS_SUBANSWER_GENERATION
+from onyx.configs.agent_configs import AGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_GENERATION
+from onyx.configs.agent_configs import AGENT_TIMEOUT_LLM_SUBANSWER_GENERATION
+from onyx.llm.chat_llm import LLMRateLimitError
+from onyx.llm.chat_llm import LLMTimeoutError
+from onyx.prompts.agent_search import NO_RECOVERED_DOCS
+from onyx.utils.logger import setup_logger
+from onyx.utils.threadpool_concurrency import run_with_timeout
+from onyx.utils.timing import log_function_time
+
+logger = setup_logger()
+
+_llm_node_error_strings = LLMNodeErrorStrings(
+    timeout="LLM Timeout Error. A sub-answer could not be constructed and the sub-question will be ignored.",
+    rate_limit="LLM Rate Limit Error. A sub-answer could not be constructed and the sub-question will be ignored.",
+    general_error="General LLM Error. A sub-answer could not be constructed and the sub-question will be ignored.",
+)
+
+
+@log_function_time(print_only=True)
+def generate_sub_answer(
+    state: AnswerQuestionState,
+    config: RunnableConfig,
+    writer: StreamWriter = lambda _: None,
+) -> SubQuestionAnswerGenerationUpdate:
+    """
+    LangGraph node to generate a sub-answer.
+    """
+    node_start_time = datetime.now()
+
+    graph_config = cast(GraphConfig, config["metadata"]["config"])
+    question = state.question
+    state.verified_reranked_documents
+    level, question_num = parse_question_id(state.question_id)
+    context_docs = state.context_documents[:AGENT_MAX_ANSWER_CONTEXT_DOCS]
+
+    context_docs = dedup_sort_inference_section_list(context_docs)
+
+    persona_contextualized_prompt = get_persona_agent_prompt_expressions(
+        graph_config.inputs.persona
+    ).contextualized_prompt
+
+    if len(context_docs) == 0:
+        answer_str = NO_RECOVERED_DOCS
+        cited_documents: list = []
+        log_results = "No documents retrieved"
+        write_custom_event(
+            "sub_answers",
+            AgentAnswerPiece(
+                answer_piece=answer_str,
+                level=level,
+                level_question_num=question_num,
+                answer_type="agent_sub_answer",
+            ),
+            writer,
+        )
+    else:
+        fast_llm = graph_config.tooling.fast_llm
+        msg = build_sub_question_answer_prompt(
+            question=question,
+            original_question=graph_config.inputs.prompt_builder.raw_user_query,
+            docs=context_docs,
+            persona_specification=persona_contextualized_prompt,
+            config=fast_llm.config,
+        )
+
+        agent_error: AgentErrorLog | None = None
+        response: list[str] = []
+
+        try:
+            response, _ = run_with_timeout(
+                AGENT_TIMEOUT_LLM_SUBANSWER_GENERATION,
+                lambda: stream_llm_answer(
+                    llm=fast_llm,
+                    prompt=msg,
+                    event_name="sub_answers",
+                    writer=writer,
+                    agent_answer_level=level,
+                    agent_answer_question_num=question_num,
+                    agent_answer_type="agent_sub_answer",
+                    timeout_override=AGENT_TIMEOUT_CONNECT_LLM_SUBANSWER_GENERATION,
+                    max_tokens=AGENT_MAX_TOKENS_SUBANSWER_GENERATION,
+                ),
+            )
+
+        except (LLMTimeoutError, TimeoutError):
+            agent_error = AgentErrorLog(
+                error_type=AgentLLMErrorType.TIMEOUT,
+                error_message=AGENT_LLM_TIMEOUT_MESSAGE,
+                error_result=_llm_node_error_strings.timeout,
+            )
+            logger.error("LLM Timeout Error - generate sub answer")
+        except LLMRateLimitError:
+            agent_error = AgentErrorLog(
+                error_type=AgentLLMErrorType.RATE_LIMIT,
+                error_message=AGENT_LLM_RATELIMIT_MESSAGE,
+                error_result=_llm_node_error_strings.rate_limit,
+            )
+            logger.error("LLM Rate Limit Error - generate sub answer")
+
+        if agent_error:
+            answer_str = LLM_ANSWER_ERROR_MESSAGE
+            cited_documents = []
+            log_results = (
+                agent_error.error_result
+                or "Sub-answer generation failed due to LLM error"
+            )
+
+        else:
+            answer_str = merge_message_runs(response, chunk_separator="")[0].content
+            answer_citation_ids = get_answer_citation_ids(answer_str)
+            cited_documents = [
+                context_docs[id] for id in answer_citation_ids if id < len(context_docs)
+            ]
+            log_results = None
+
+    stop_event = StreamStopInfo(
+        stop_reason=StreamStopReason.FINISHED,
+        stream_type=StreamType.SUB_ANSWER,
+        level=level,
+        level_question_num=question_num,
+    )
+    write_custom_event("stream_finished", stop_event, writer)
+
+    return SubQuestionAnswerGenerationUpdate(
+        answer=answer_str,
+        cited_documents=cited_documents,
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="initial - generate individual sub answer",
+                node_name="generate sub answer",
+                node_start_time=node_start_time,
+                result=log_results or "",
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/nodes/ingest_retrieved_documents.py

@@ -0,0 +1,25 @@
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionRetrievalIngestionUpdate,
+)
+from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkRetrievalStats
+
+
+def ingest_retrieved_documents(
+    state: ExpandedRetrievalOutput,
+) -> SubQuestionRetrievalIngestionUpdate:
+    """
+    LangGraph node to ingest the retrieved documents to format it for the sub-answer.
+    """
+    sub_question_retrieval_stats = state.expanded_retrieval_result.retrieval_stats
+    if sub_question_retrieval_stats is None:
+        sub_question_retrieval_stats = [AgentChunkRetrievalStats()]
+
+    return SubQuestionRetrievalIngestionUpdate(
+        expanded_retrieval_results=state.expanded_retrieval_result.expanded_query_results,
+        verified_reranked_documents=state.expanded_retrieval_result.verified_reranked_documents,
+        context_documents=state.expanded_retrieval_result.context_documents,
+        sub_question_retrieval_stats=sub_question_retrieval_stats,
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/states.py

@@ -0,0 +1,73 @@
+from operator import add
+from typing import Annotated
+
+from pydantic import BaseModel
+
+from onyx.agents.agent_search.core_state import SubgraphCoreState
+from onyx.agents.agent_search.deep_search.main.states import LoggerUpdate
+from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkRetrievalStats
+from onyx.agents.agent_search.shared_graph_utils.models import QueryRetrievalResult
+from onyx.agents.agent_search.shared_graph_utils.models import (
+    SubQuestionAnswerResults,
+)
+from onyx.agents.agent_search.shared_graph_utils.operators import (
+    dedup_inference_sections,
+)
+from onyx.context.search.models import InferenceSection
+
+
+## Update States
+class SubQuestionAnswerCheckUpdate(LoggerUpdate, BaseModel):
+    answer_quality: bool = False
+    log_messages: list[str] = []
+
+
+class SubQuestionAnswerGenerationUpdate(LoggerUpdate, BaseModel):
+    answer: str = ""
+    log_messages: list[str] = []
+    cited_documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
+    # answer_stat: AnswerStats
+
+
+class SubQuestionRetrievalIngestionUpdate(LoggerUpdate, BaseModel):
+    expanded_retrieval_results: list[QueryRetrievalResult] = []
+    verified_reranked_documents: Annotated[
+        list[InferenceSection], dedup_inference_sections
+    ] = []
+    context_documents: Annotated[list[InferenceSection], dedup_inference_sections] = []
+    sub_question_retrieval_stats: AgentChunkRetrievalStats = AgentChunkRetrievalStats()
+
+
+## Graph Input State
+
+
+class SubQuestionAnsweringInput(SubgraphCoreState):
+    question: str
+    question_id: str
+    # level 0 is original question and first decomposition, level 1 is follow up, etc
+    # question_num is a unique number per original question per level.
+
+
+## Graph State
+
+
+class AnswerQuestionState(
+    SubQuestionAnsweringInput,
+    SubQuestionAnswerGenerationUpdate,
+    SubQuestionAnswerCheckUpdate,
+    SubQuestionRetrievalIngestionUpdate,
+):
+    pass
+
+
+## Graph Output State
+
+
+class AnswerQuestionOutput(LoggerUpdate, BaseModel):
+    """
+    This is a list of results even though each call of this subgraph only returns one result.
+    This is because if we parallelize the answer query subgraph, there will be multiple
+      results in a list so the add operator is used to add them together.
+    """
+
+    answer_results: Annotated[list[SubQuestionAnswerResults], add] = []

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/edges.py

@@ -0,0 +1,50 @@
+from collections.abc import Hashable
+from datetime import datetime
+
+from langgraph.types import Send
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionAnsweringInput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalState,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import make_question_id
+
+
+def parallelize_initial_sub_question_answering(
+    state: SubQuestionRetrievalState,
+) -> list[Send | Hashable]:
+    """
+    LangGraph edge to parallelize the initial sub-question answering. If there are no sub-questions,
+    we send empty answers to the initial answer generation, and that answer would be generated
+    solely based on the documents retrieved for the original question.
+    """
+    edge_start_time = datetime.now()
+    if len(state.initial_sub_questions) > 0:
+        return [
+            Send(
+                "answer_query_subgraph",
+                SubQuestionAnsweringInput(
+                    question=question,
+                    question_id=make_question_id(0, question_num + 1),
+                    log_messages=[
+                        f"{edge_start_time} -- Main Edge - Parallelize Initial Sub-question Answering"
+                    ],
+                ),
+            )
+            for question_num, question in enumerate(state.initial_sub_questions)
+        ]
+
+    else:
+        return [
+            Send(
+                "ingest_answers",
+                AnswerQuestionOutput(
+                    answer_results=[],
+                ),
+            )
+        ]

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/graph_builder.py

@@ -0,0 +1,96 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.nodes.generate_initial_answer import (
+    generate_initial_answer,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.nodes.validate_initial_answer import (
+    validate_initial_answer,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalInput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalState,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_sub_answers.graph_builder import (
+    generate_sub_answers_graph_builder,
+)
+from onyx.agents.agent_search.deep_search.initial.retrieve_orig_question_docs.graph_builder import (
+    retrieve_orig_question_docs_graph_builder,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def generate_initial_answer_graph_builder(test_mode: bool = False) -> StateGraph:
+    """
+    LangGraph graph builder for the initial answer generation.
+    """
+    graph = StateGraph(
+        state_schema=SubQuestionRetrievalState,
+        input=SubQuestionRetrievalInput,
+    )
+
+    # The sub-graph that generates the initial sub-answers
+    generate_sub_answers = generate_sub_answers_graph_builder().compile()
+    graph.add_node(
+        node="generate_sub_answers_subgraph",
+        action=generate_sub_answers,
+    )
+
+    # The sub-graph that retrieves the original question documents. This is run
+    # in parallel with the sub-answer generation process
+    retrieve_orig_question_docs = retrieve_orig_question_docs_graph_builder().compile()
+    graph.add_node(
+        node="retrieve_orig_question_docs_subgraph_wrapper",
+        action=retrieve_orig_question_docs,
+    )
+
+    # Node that generates the initial answer using the results of the previous
+    # two sub-graphs
+    graph.add_node(
+        node="generate_initial_answer",
+        action=generate_initial_answer,
+    )
+
+    # Node that validates the initial answer
+    graph.add_node(
+        node="validate_initial_answer",
+        action=validate_initial_answer,
+    )
+
+    ### Add edges ###
+
+    graph.add_edge(
+        start_key=START,
+        end_key="retrieve_orig_question_docs_subgraph_wrapper",
+    )
+
+    graph.add_edge(
+        start_key=START,
+        end_key="generate_sub_answers_subgraph",
+    )
+
+    # Wait for both, the original question docs and the sub-answers to be generated before proceeding
+    graph.add_edge(
+        start_key=[
+            "retrieve_orig_question_docs_subgraph_wrapper",
+            "generate_sub_answers_subgraph",
+        ],
+        end_key="generate_initial_answer",
+    )
+
+    graph.add_edge(
+        start_key="generate_initial_answer",
+        end_key="validate_initial_answer",
+    )
+
+    graph.add_edge(
+        start_key="validate_initial_answer",
+        end_key=END,
+    )
+
+    return graph

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/nodes/generate_initial_answer.py

@@ -0,0 +1,405 @@
+from datetime import datetime
+from typing import cast
+
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import merge_content
+from langchain_core.runnables import RunnableConfig
+from langgraph.types import StreamWriter
+
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalState,
+)
+from onyx.agents.agent_search.deep_search.main.models import AgentBaseMetrics
+from onyx.agents.agent_search.deep_search.main.operations import (
+    calculate_initial_agent_stats,
+)
+from onyx.agents.agent_search.deep_search.main.operations import get_query_info
+from onyx.agents.agent_search.deep_search.main.operations import logger
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialAnswerUpdate,
+)
+from onyx.agents.agent_search.models import GraphConfig
+from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
+    get_prompt_enrichment_components,
+)
+from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
+    trim_prompt_piece,
+)
+from onyx.agents.agent_search.shared_graph_utils.calculations import (
+    get_answer_generation_documents,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_LLM_RATELIMIT_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AGENT_LLM_TIMEOUT_MESSAGE,
+)
+from onyx.agents.agent_search.shared_graph_utils.constants import (
+    AgentLLMErrorType,
+)
+from onyx.agents.agent_search.shared_graph_utils.llm import stream_llm_answer
+from onyx.agents.agent_search.shared_graph_utils.models import AgentErrorLog
+from onyx.agents.agent_search.shared_graph_utils.models import InitialAgentResultStats
+from onyx.agents.agent_search.shared_graph_utils.models import LLMNodeErrorStrings
+from onyx.agents.agent_search.shared_graph_utils.operators import (
+    dedup_inference_section_list,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import _should_restrict_tokens
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    dispatch_main_answer_stop_info,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import format_docs
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_deduplicated_structured_subquestion_documents,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_langgraph_node_log_string,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import relevance_from_docs
+from onyx.agents.agent_search.shared_graph_utils.utils import remove_document_citations
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import AgentAnswerPiece
+from onyx.chat.models import ExtendedToolResponse
+from onyx.chat.models import StreamingError
+from onyx.configs.agent_configs import AGENT_ANSWER_GENERATION_BY_FAST_LLM
+from onyx.configs.agent_configs import AGENT_MAX_ANSWER_CONTEXT_DOCS
+from onyx.configs.agent_configs import AGENT_MAX_STREAMED_DOCS_FOR_INITIAL_ANSWER
+from onyx.configs.agent_configs import AGENT_MAX_TOKENS_ANSWER_GENERATION
+from onyx.configs.agent_configs import AGENT_MIN_ORIG_QUESTION_DOCS
+from onyx.configs.agent_configs import (
+    AGENT_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION,
+)
+from onyx.configs.agent_configs import (
+    AGENT_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION,
+)
+from onyx.llm.chat_llm import LLMRateLimitError
+from onyx.llm.chat_llm import LLMTimeoutError
+from onyx.prompts.agent_search import INITIAL_ANSWER_PROMPT_W_SUB_QUESTIONS
+from onyx.prompts.agent_search import (
+    INITIAL_ANSWER_PROMPT_WO_SUB_QUESTIONS,
+)
+from onyx.prompts.agent_search import (
+    SUB_QUESTION_ANSWER_TEMPLATE,
+)
+from onyx.prompts.agent_search import UNKNOWN_ANSWER
+from onyx.tools.tool_implementations.search.search_tool import yield_search_responses
+from onyx.utils.threadpool_concurrency import run_with_timeout
+from onyx.utils.timing import log_function_time
+
+_llm_node_error_strings = LLMNodeErrorStrings(
+    timeout="LLM Timeout Error. The initial answer could not be generated.",
+    rate_limit="LLM Rate Limit Error. The initial answer could not be generated.",
+    general_error="General LLM Error. The initial answer could not be generated.",
+)
+
+
+@log_function_time(print_only=True)
+def generate_initial_answer(
+    state: SubQuestionRetrievalState,
+    config: RunnableConfig,
+    writer: StreamWriter = lambda _: None,
+) -> InitialAnswerUpdate:
+    """
+    LangGraph node to generate the initial answer, using the initial sub-questions/sub-answers and the
+    documents retrieved for the original question.
+    """
+    node_start_time = datetime.now()
+
+    graph_config = cast(GraphConfig, config["metadata"]["config"])
+    question = graph_config.inputs.prompt_builder.raw_user_query
+    prompt_enrichment_components = get_prompt_enrichment_components(graph_config)
+
+    # get all documents cited in sub-questions
+    structured_subquestion_docs = get_deduplicated_structured_subquestion_documents(
+        state.sub_question_results
+    )
+
+    orig_question_retrieval_documents = state.orig_question_retrieved_documents
+
+    consolidated_context_docs = structured_subquestion_docs.cited_documents
+    counter = 0
+    for original_doc in orig_question_retrieval_documents:
+        if original_doc in structured_subquestion_docs.cited_documents:
+            continue
+
+        if (
+            counter <= AGENT_MIN_ORIG_QUESTION_DOCS
+            or len(consolidated_context_docs) < AGENT_MAX_ANSWER_CONTEXT_DOCS
+        ):
+            consolidated_context_docs.append(original_doc)
+            counter += 1
+
+    # sort docs by their scores - though the scores refer to different questions
+    relevant_docs = dedup_inference_section_list(consolidated_context_docs)
+
+    sub_questions: list[str] = []
+
+    # Create the list of documents to stream out. Start with the
+    # ones that wil be in the context (or, if len == 0, use docs
+    # that were retrieved for the original question)
+    answer_generation_documents = get_answer_generation_documents(
+        relevant_docs=relevant_docs,
+        context_documents=structured_subquestion_docs.context_documents,
+        original_question_docs=orig_question_retrieval_documents,
+        max_docs=AGENT_MAX_STREAMED_DOCS_FOR_INITIAL_ANSWER,
+    )
+
+    # Use the query info from the base document retrieval
+    query_info = get_query_info(state.orig_question_sub_query_retrieval_results)
+
+    assert (
+        graph_config.tooling.search_tool
+    ), "search_tool must be provided for agentic search"
+
+    relevance_list = relevance_from_docs(
+        answer_generation_documents.streaming_documents
+    )
+    for tool_response in yield_search_responses(
+        query=question,
+        get_retrieved_sections=lambda: answer_generation_documents.context_documents,
+        get_final_context_sections=lambda: answer_generation_documents.context_documents,
+        search_query_info=query_info,
+        get_section_relevance=lambda: relevance_list,
+        search_tool=graph_config.tooling.search_tool,
+    ):
+        write_custom_event(
+            "tool_response",
+            ExtendedToolResponse(
+                id=tool_response.id,
+                response=tool_response.response,
+                level=0,
+                level_question_num=0,  # 0, 0 is the base question
+            ),
+            writer,
+        )
+
+    if len(answer_generation_documents.context_documents) == 0:
+        write_custom_event(
+            "initial_agent_answer",
+            AgentAnswerPiece(
+                answer_piece=UNKNOWN_ANSWER,
+                level=0,
+                level_question_num=0,
+                answer_type="agent_level_answer",
+            ),
+            writer,
+        )
+        dispatch_main_answer_stop_info(0, writer)
+
+        answer = UNKNOWN_ANSWER
+        initial_agent_stats = InitialAgentResultStats(
+            sub_questions={},
+            original_question={},
+            agent_effectiveness={},
+        )
+
+    else:
+        sub_question_answer_results = state.sub_question_results
+
+        # Collect the sub-questions and sub-answers and construct an appropriate
+        # prompt string.
+        # Consider replacing by a function.
+        answered_sub_questions: list[str] = []
+        all_sub_questions: list[str] = []  # Separate list for tracking all questions
+
+        for idx, sub_question_answer_result in enumerate(
+            sub_question_answer_results, start=1
+        ):
+            all_sub_questions.append(sub_question_answer_result.question)
+
+            is_valid_answer = (
+                sub_question_answer_result.verified_high_quality
+                and sub_question_answer_result.answer
+                and sub_question_answer_result.answer != UNKNOWN_ANSWER
+            )
+
+            if is_valid_answer:
+                answered_sub_questions.append(
+                    SUB_QUESTION_ANSWER_TEMPLATE.format(
+                        sub_question=sub_question_answer_result.question,
+                        sub_answer=sub_question_answer_result.answer,
+                        sub_question_num=idx,
+                    )
+                )
+
+        sub_question_answer_str = (
+            "\n\n------\n\n".join(answered_sub_questions)
+            if answered_sub_questions
+            else ""
+        )
+
+        # Use the appropriate prompt based on whether there are sub-questions.
+        base_prompt = (
+            INITIAL_ANSWER_PROMPT_W_SUB_QUESTIONS
+            if answered_sub_questions
+            else INITIAL_ANSWER_PROMPT_WO_SUB_QUESTIONS
+        )
+
+        sub_questions = all_sub_questions  # Replace the original assignment
+
+        model = (
+            graph_config.tooling.fast_llm
+            if AGENT_ANSWER_GENERATION_BY_FAST_LLM
+            else graph_config.tooling.primary_llm
+        )
+
+        doc_context = format_docs(answer_generation_documents.context_documents)
+        doc_context = trim_prompt_piece(
+            config=model.config,
+            prompt_piece=doc_context,
+            reserved_str=(
+                base_prompt
+                + sub_question_answer_str
+                + prompt_enrichment_components.persona_prompts.contextualized_prompt
+                + prompt_enrichment_components.history
+                + prompt_enrichment_components.date_str
+            ),
+        )
+
+        msg = [
+            HumanMessage(
+                content=base_prompt.format(
+                    question=question,
+                    answered_sub_questions=remove_document_citations(
+                        sub_question_answer_str
+                    ),
+                    relevant_docs=doc_context,
+                    persona_specification=prompt_enrichment_components.persona_prompts.contextualized_prompt,
+                    history=prompt_enrichment_components.history,
+                    date_prompt=prompt_enrichment_components.date_str,
+                )
+            )
+        ]
+
+        streamed_tokens: list[str] = [""]
+        dispatch_timings: list[float] = []
+
+        agent_error: AgentErrorLog | None = None
+
+        try:
+            streamed_tokens, dispatch_timings = run_with_timeout(
+                AGENT_TIMEOUT_LLM_INITIAL_ANSWER_GENERATION,
+                lambda: stream_llm_answer(
+                    llm=model,
+                    prompt=msg,
+                    event_name="initial_agent_answer",
+                    writer=writer,
+                    agent_answer_level=0,
+                    agent_answer_question_num=0,
+                    agent_answer_type="agent_level_answer",
+                    timeout_override=AGENT_TIMEOUT_CONNECT_LLM_INITIAL_ANSWER_GENERATION,
+                    max_tokens=(
+                        AGENT_MAX_TOKENS_ANSWER_GENERATION
+                        if _should_restrict_tokens(model.config)
+                        else None
+                    ),
+                ),
+            )
+
+        except (LLMTimeoutError, TimeoutError):
+            agent_error = AgentErrorLog(
+                error_type=AgentLLMErrorType.TIMEOUT,
+                error_message=AGENT_LLM_TIMEOUT_MESSAGE,
+                error_result=_llm_node_error_strings.timeout,
+            )
+            logger.error("LLM Timeout Error - generate initial answer")
+
+        except LLMRateLimitError:
+            agent_error = AgentErrorLog(
+                error_type=AgentLLMErrorType.RATE_LIMIT,
+                error_message=AGENT_LLM_RATELIMIT_MESSAGE,
+                error_result=_llm_node_error_strings.rate_limit,
+            )
+            logger.error("LLM Rate Limit Error - generate initial answer")
+
+        if agent_error:
+            write_custom_event(
+                "initial_agent_answer",
+                StreamingError(
+                    error=AGENT_LLM_TIMEOUT_MESSAGE,
+                ),
+                writer,
+            )
+            return InitialAnswerUpdate(
+                initial_answer=None,
+                answer_error=AgentErrorLog(
+                    error_message=agent_error.error_message or "An LLM error occurred",
+                    error_type=agent_error.error_type,
+                    error_result=agent_error.error_result,
+                ),
+                initial_agent_stats=None,
+                generated_sub_questions=sub_questions,
+                agent_base_end_time=None,
+                agent_base_metrics=None,
+                log_messages=[
+                    get_langgraph_node_log_string(
+                        graph_component="initial - generate initial answer",
+                        node_name="generate initial answer",
+                        node_start_time=node_start_time,
+                        result=agent_error.error_result or "An LLM error occurred",
+                    )
+                ],
+            )
+
+        logger.debug(
+            f"Average dispatch time for initial answer: {sum(dispatch_timings) / len(dispatch_timings)}"
+        )
+
+        dispatch_main_answer_stop_info(0, writer)
+        response = merge_content(*streamed_tokens)
+        answer = cast(str, response)
+
+        initial_agent_stats = calculate_initial_agent_stats(
+            state.sub_question_results, state.orig_question_retrieval_stats
+        )
+
+        logger.debug(
+            f"\n\nYYYYY--Sub-Questions:\n\n{sub_question_answer_str}\n\nStats:\n\n"
+        )
+
+        if initial_agent_stats:
+            logger.debug(initial_agent_stats.original_question)
+            logger.debug(initial_agent_stats.sub_questions)
+            logger.debug(initial_agent_stats.agent_effectiveness)
+
+    agent_base_end_time = datetime.now()
+
+    if agent_base_end_time and state.agent_start_time:
+        duration_s = (agent_base_end_time - state.agent_start_time).total_seconds()
+    else:
+        duration_s = None
+
+    agent_base_metrics = AgentBaseMetrics(
+        num_verified_documents_total=len(relevant_docs),
+        num_verified_documents_core=state.orig_question_retrieval_stats.verified_count,
+        verified_avg_score_core=state.orig_question_retrieval_stats.verified_avg_scores,
+        num_verified_documents_base=initial_agent_stats.sub_questions.get(
+            "num_verified_documents"
+        ),
+        verified_avg_score_base=initial_agent_stats.sub_questions.get(
+            "verified_avg_score"
+        ),
+        base_doc_boost_factor=initial_agent_stats.agent_effectiveness.get(
+            "utilized_chunk_ratio"
+        ),
+        support_boost_factor=initial_agent_stats.agent_effectiveness.get(
+            "support_ratio"
+        ),
+        duration_s=duration_s,
+    )
+
+    return InitialAnswerUpdate(
+        initial_answer=answer,
+        initial_agent_stats=initial_agent_stats,
+        generated_sub_questions=sub_questions,
+        agent_base_end_time=agent_base_end_time,
+        agent_base_metrics=agent_base_metrics,
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="initial - generate initial answer",
+                node_name="generate initial answer",
+                node_start_time=node_start_time,
+                result="",
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/nodes/validate_initial_answer.py

@@ -0,0 +1,42 @@
+from datetime import datetime
+
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalState,
+)
+from onyx.agents.agent_search.deep_search.main.operations import logger
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialAnswerQualityUpdate,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_langgraph_node_log_string,
+)
+from onyx.utils.timing import log_function_time
+
+
+@log_function_time(print_only=True)
+def validate_initial_answer(
+    state: SubQuestionRetrievalState,
+) -> InitialAnswerQualityUpdate:
+    """
+    Check whether the initial answer sufficiently addresses the original user question.
+    """
+
+    node_start_time = datetime.now()
+
+    logger.debug(
+        f"--------{node_start_time}--------Checking for base answer validity - for not set True/False manually"
+    )
+
+    verdict = True  # not actually required as already streamed out. Refinement will do similar
+
+    return InitialAnswerQualityUpdate(
+        initial_answer_quality_eval=verdict,
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="initial - generate initial answer",
+                node_name="validate initial answer",
+                node_start_time=node_start_time,
+                result="",
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_initial_answer/states.py

@@ -0,0 +1,51 @@
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agents.agent_search.core_state import CoreState
+from onyx.agents.agent_search.deep_search.main.states import (
+    ExploratorySearchUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialAnswerQualityUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialAnswerUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialQuestionDecompositionUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    OrigQuestionRetrievalUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    SubQuestionResultsUpdate,
+)
+from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.models import (
+    QuestionRetrievalResult,
+)
+from onyx.context.search.models import InferenceSection
+
+
+### States ###
+class SubQuestionRetrievalInput(CoreState):
+    exploratory_search_results: list[InferenceSection]
+
+
+## Graph State
+class SubQuestionRetrievalState(
+    # This includes the core state
+    SubQuestionRetrievalInput,
+    InitialQuestionDecompositionUpdate,
+    InitialAnswerUpdate,
+    SubQuestionResultsUpdate,
+    OrigQuestionRetrievalUpdate,
+    InitialAnswerQualityUpdate,
+    ExploratorySearchUpdate,
+):
+    base_raw_search_result: Annotated[list[QuestionRetrievalResult], add]
+
+
+## Graph Output State
+class SubQuestionRetrievalOutput(TypedDict):
+    log_messages: list[str]

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/edges.py

@@ -0,0 +1,48 @@
+from collections.abc import Hashable
+from datetime import datetime
+
+from langgraph.types import Send
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    SubQuestionAnsweringInput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalState,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import make_question_id
+
+
+def parallelize_initial_sub_question_answering(
+    state: SubQuestionRetrievalState,
+) -> list[Send | Hashable]:
+    """
+    LangGraph edge to parallelize the initial sub-question answering.
+    """
+    edge_start_time = datetime.now()
+    if len(state.initial_sub_questions) > 0:
+        return [
+            Send(
+                "answer_sub_question_subgraphs",
+                SubQuestionAnsweringInput(
+                    question=question,
+                    question_id=make_question_id(0, question_num + 1),
+                    log_messages=[
+                        f"{edge_start_time} -- Main Edge - Parallelize Initial Sub-question Answering"
+                    ],
+                ),
+            )
+            for question_num, question in enumerate(state.initial_sub_questions)
+        ]
+
+    else:
+        return [
+            Send(
+                "ingest_answers",
+                AnswerQuestionOutput(
+                    answer_results=[],
+                ),
+            )
+        ]

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/graph_builder.py

@@ -0,0 +1,81 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.graph_builder import (
+    answer_query_graph_builder,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_sub_answers.edges import (
+    parallelize_initial_sub_question_answering,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_sub_answers.nodes.decompose_orig_question import (
+    decompose_orig_question,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_sub_answers.nodes.format_initial_sub_answers import (
+    format_initial_sub_answers,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_sub_answers.states import (
+    SubQuestionAnsweringInput,
+)
+from onyx.agents.agent_search.deep_search.initial.generate_sub_answers.states import (
+    SubQuestionAnsweringState,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+test_mode = False
+
+
+def generate_sub_answers_graph_builder() -> StateGraph:
+    """
+    LangGraph graph builder for the initial sub-answer generation process.
+    It generates the initial sub-questions and produces the answers.
+    """
+
+    graph = StateGraph(
+        state_schema=SubQuestionAnsweringState,
+        input=SubQuestionAnsweringInput,
+    )
+
+    # Decompose the original question into sub-questions
+    graph.add_node(
+        node="decompose_orig_question",
+        action=decompose_orig_question,
+    )
+
+    # The sub-graph that executes the initial sub-question answering for
+    # each of the sub-questions.
+    answer_sub_question_subgraphs = answer_query_graph_builder().compile()
+    graph.add_node(
+        node="answer_sub_question_subgraphs",
+        action=answer_sub_question_subgraphs,
+    )
+
+    # Node that collects and formats the initial sub-question answers
+    graph.add_node(
+        node="format_initial_sub_question_answers",
+        action=format_initial_sub_answers,
+    )
+
+    graph.add_edge(
+        start_key=START,
+        end_key="decompose_orig_question",
+    )
+
+    graph.add_conditional_edges(
+        source="decompose_orig_question",
+        path=parallelize_initial_sub_question_answering,
+        path_map=["answer_sub_question_subgraphs"],
+    )
+    graph.add_edge(
+        start_key=["answer_sub_question_subgraphs"],
+        end_key="format_initial_sub_question_answers",
+    )
+
+    graph.add_edge(
+        start_key="format_initial_sub_question_answers",
+        end_key=END,
+    )
+
+    return graph

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/nodes/decompose_orig_question.py

@@ -0,0 +1,190 @@
+from datetime import datetime
+from typing import cast
+
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import merge_content
+from langchain_core.runnables import RunnableConfig
+from langgraph.types import StreamWriter
+
+from onyx.agents.agent_search.deep_search.initial.generate_initial_answer.states import (
+    SubQuestionRetrievalState,
+)
+from onyx.agents.agent_search.deep_search.main.models import (
+    AgentRefinedMetrics,
+)
+from onyx.agents.agent_search.deep_search.main.operations import dispatch_subquestion
+from onyx.agents.agent_search.deep_search.main.operations import (
+    dispatch_subquestion_sep,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialQuestionDecompositionUpdate,
+)
+from onyx.agents.agent_search.models import GraphConfig
+from onyx.agents.agent_search.shared_graph_utils.agent_prompt_ops import (
+    build_history_prompt,
+)
+from onyx.agents.agent_search.shared_graph_utils.models import BaseMessage_Content
+from onyx.agents.agent_search.shared_graph_utils.models import LLMNodeErrorStrings
+from onyx.agents.agent_search.shared_graph_utils.utils import dispatch_separated
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_langgraph_node_log_string,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event
+from onyx.chat.models import StreamStopInfo
+from onyx.chat.models import StreamStopReason
+from onyx.chat.models import StreamType
+from onyx.chat.models import SubQuestionPiece
+from onyx.configs.agent_configs import AGENT_MAX_TOKENS_SUBQUESTION_GENERATION
+from onyx.configs.agent_configs import AGENT_NUM_DOCS_FOR_DECOMPOSITION
+from onyx.configs.agent_configs import (
+    AGENT_TIMEOUT_CONNECT_LLM_SUBQUESTION_GENERATION,
+)
+from onyx.configs.agent_configs import (
+    AGENT_TIMEOUT_LLM_SUBQUESTION_GENERATION,
+)
+from onyx.llm.chat_llm import LLMRateLimitError
+from onyx.llm.chat_llm import LLMTimeoutError
+from onyx.prompts.agent_search import (
+    INITIAL_DECOMPOSITION_PROMPT_QUESTIONS_AFTER_SEARCH_ASSUMING_REFINEMENT,
+)
+from onyx.prompts.agent_search import (
+    INITIAL_QUESTION_DECOMPOSITION_PROMPT_ASSUMING_REFINEMENT,
+)
+from onyx.utils.logger import setup_logger
+from onyx.utils.threadpool_concurrency import run_with_timeout
+from onyx.utils.timing import log_function_time
+
+logger = setup_logger()
+
+_llm_node_error_strings = LLMNodeErrorStrings(
+    timeout="LLM Timeout Error. Sub-questions could not be generated.",
+    rate_limit="LLM Rate Limit Error. Sub-questions could not be generated.",
+    general_error="General LLM Error. Sub-questions could not be generated.",
+)
+
+
+@log_function_time(print_only=True)
+def decompose_orig_question(
+    state: SubQuestionRetrievalState,
+    config: RunnableConfig,
+    writer: StreamWriter = lambda _: None,
+) -> InitialQuestionDecompositionUpdate:
+    """
+    LangGraph node to decompose the original question into sub-questions.
+    """
+    node_start_time = datetime.now()
+
+    graph_config = cast(GraphConfig, config["metadata"]["config"])
+    question = graph_config.inputs.prompt_builder.raw_user_query
+    perform_initial_search_decomposition = (
+        graph_config.behavior.perform_initial_search_decomposition
+    )
+    # Get the rewritten queries in a defined format
+    model = graph_config.tooling.fast_llm
+
+    history = build_history_prompt(graph_config, question)
+
+    # Use the initial search results to inform the decomposition
+    agent_start_time = datetime.now()
+
+    # Initial search to inform decomposition. Just get top 3 fits
+
+    if perform_initial_search_decomposition:
+        # Due to unfortunate state representation in LangGraph, we need here to double check that the retrieval has
+        # happened prior to this point, allowing silent failure here since it is not critical for decomposition in
+        # all queries.
+        if not state.exploratory_search_results:
+            logger.error("Initial search for decomposition failed")
+
+        sample_doc_str = "\n\n".join(
+            [
+                doc.combined_content
+                for doc in state.exploratory_search_results[
+                    :AGENT_NUM_DOCS_FOR_DECOMPOSITION
+                ]
+            ]
+        )
+
+        decomposition_prompt = INITIAL_DECOMPOSITION_PROMPT_QUESTIONS_AFTER_SEARCH_ASSUMING_REFINEMENT.format(
+            question=question, sample_doc_str=sample_doc_str, history=history
+        )
+
+    else:
+        decomposition_prompt = (
+            INITIAL_QUESTION_DECOMPOSITION_PROMPT_ASSUMING_REFINEMENT.format(
+                question=question, history=history
+            )
+        )
+
+    # Start decomposition
+
+    msg = [HumanMessage(content=decomposition_prompt)]
+
+    # Send the initial question as a subquestion with number 0
+    write_custom_event(
+        "decomp_qs",
+        SubQuestionPiece(
+            sub_question=question,
+            level=0,
+            level_question_num=0,
+        ),
+        writer,
+    )
+
+    # dispatches custom events for subquestion tokens, adding in subquestion ids.
+
+    streamed_tokens: list[BaseMessage_Content] = []
+
+    try:
+        streamed_tokens = run_with_timeout(
+            AGENT_TIMEOUT_LLM_SUBQUESTION_GENERATION,
+            dispatch_separated,
+            model.stream(
+                msg,
+                timeout_override=AGENT_TIMEOUT_CONNECT_LLM_SUBQUESTION_GENERATION,
+                max_tokens=AGENT_MAX_TOKENS_SUBQUESTION_GENERATION,
+            ),
+            dispatch_subquestion(0, writer),
+            sep_callback=dispatch_subquestion_sep(0, writer),
+        )
+
+        decomposition_response = merge_content(*streamed_tokens)
+
+        list_of_subqs = cast(str, decomposition_response).split("\n")
+
+        initial_sub_questions = [sq.strip() for sq in list_of_subqs if sq.strip() != ""]
+        log_result = f"decomposed original question into {len(initial_sub_questions)} subquestions"
+
+        stop_event = StreamStopInfo(
+            stop_reason=StreamStopReason.FINISHED,
+            stream_type=StreamType.SUB_QUESTIONS,
+            level=0,
+        )
+        write_custom_event("stream_finished", stop_event, writer)
+
+    except (LLMTimeoutError, TimeoutError) as e:
+        logger.error("LLM Timeout Error - decompose orig question")
+        raise e  # fail loudly on this critical step
+    except LLMRateLimitError as e:
+        logger.error("LLM Rate Limit Error - decompose orig question")
+        raise e
+
+    return InitialQuestionDecompositionUpdate(
+        initial_sub_questions=initial_sub_questions,
+        agent_start_time=agent_start_time,
+        agent_refined_start_time=None,
+        agent_refined_end_time=None,
+        agent_refined_metrics=AgentRefinedMetrics(
+            refined_doc_boost_factor=None,
+            refined_question_boost_factor=None,
+            duration_s=None,
+        ),
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="initial - generate sub answers",
+                node_name="decompose original question",
+                node_start_time=node_start_time,
+                result=log_result,
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/nodes/format_initial_sub_answers.py

@@ -0,0 +1,50 @@
+from datetime import datetime
+
+from onyx.agents.agent_search.deep_search.initial.generate_individual_sub_answer.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    SubQuestionResultsUpdate,
+)
+from onyx.agents.agent_search.shared_graph_utils.operators import (
+    dedup_inference_sections,
+)
+from onyx.agents.agent_search.shared_graph_utils.utils import (
+    get_langgraph_node_log_string,
+)
+
+
+def format_initial_sub_answers(
+    state: AnswerQuestionOutput,
+) -> SubQuestionResultsUpdate:
+    """
+    LangGraph node to format the answers to the initial sub-questions, including
+    deduping verified documents and context documents.
+    """
+    node_start_time = datetime.now()
+
+    documents = []
+    context_documents = []
+    cited_documents = []
+    answer_results = state.answer_results
+    for answer_result in answer_results:
+        documents.extend(answer_result.verified_reranked_documents)
+        context_documents.extend(answer_result.context_documents)
+        cited_documents.extend(answer_result.cited_documents)
+
+    return SubQuestionResultsUpdate(
+        # Deduping is done by the documents operator for the main graph
+        # so we might not need to dedup here
+        verified_reranked_documents=dedup_inference_sections(documents, []),
+        context_documents=dedup_inference_sections(context_documents, []),
+        cited_documents=dedup_inference_sections(cited_documents, []),
+        sub_question_results=answer_results,
+        log_messages=[
+            get_langgraph_node_log_string(
+                graph_component="initial - generate sub answers",
+                node_name="format initial sub answers",
+                node_start_time=node_start_time,
+                result="",
+            )
+        ],
+    )

backend/onyx/agents/agent_search/deep_search/initial/generate_sub_answers/states.py

@@ -0,0 +1,34 @@
+from typing import TypedDict
+
+from onyx.agents.agent_search.core_state import CoreState
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialAnswerUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    InitialQuestionDecompositionUpdate,
+)
+from onyx.agents.agent_search.deep_search.main.states import (
+    SubQuestionResultsUpdate,
+)
+from onyx.context.search.models import InferenceSection
+
+
+### States ###
+class SubQuestionAnsweringInput(CoreState):
+    exploratory_search_results: list[InferenceSection]
+
+
+## Graph State
+class SubQuestionAnsweringState(
+    # This includes the core state
+    SubQuestionAnsweringInput,
+    InitialQuestionDecompositionUpdate,
+    InitialAnswerUpdate,
+    SubQuestionResultsUpdate,
+):
+    pass
+
+
+## Graph Output State
+class SubQuestionAnsweringOutput(TypedDict):
+    log_messages: list[str]