agents

Co-authored-by: cubic-dev-ai[bot] <191113872+cubic-dev-ai[bot]@users.noreply.github.com>

.github/CODEOWNERS

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

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

@@ -0,0 +1,49 @@
+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,6 +13,14 @@ 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,7 +55,25 @@ 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" --debug --config ct.yaml
+      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
       # 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,6 +19,10 @@ 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:
@@ -66,7 +70,9 @@ jobs:
             -i /local/openapi.json \
             -g python \
             -o /local/onyx_openapi_client \
-            --package-name onyx_openapi_client
+            --package-name onyx_openapi_client \
+            --skip-validate-spec \
+            --openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -248,6 +254,8 @@ 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..."
@@ -272,6 +280,10 @@ 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

@@ -0,0 +1,38 @@
+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,6 +19,10 @@ 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:
@@ -63,7 +67,9 @@ jobs:
             -i /local/openapi.json \
             -g python \
             -o /local/onyx_openapi_client \
-            --package-name onyx_openapi_client
+            --package-name onyx_openapi_client \
+            --skip-validate-spec \
+            --openapi-normalizer "SIMPLIFY_ONEOF_ANYOF=true,SET_OAS3_NULLABLE=true"
             
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -207,6 +213,10 @@ 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-playwright-tests.yml

@@ -10,7 +10,6 @@ env:
   SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
   GEN_AI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
   MOCK_LLM_RESPONSE: true
-  PYTEST_PLAYWRIGHT_SKIP_INITIAL_RESET: true
 
 jobs:
   playwright-tests:
@@ -160,12 +159,6 @@ jobs:
           done
           echo "Finished waiting for service."
 
-      - name: Run pytest playwright test init
-        working-directory: ./backend
-        env:
-          PYTEST_IGNORE_SKIP: true
-        run: pytest -s tests/integration/tests/playwright/test_playwright.py
-
       - name: Run Playwright tests
         working-directory: ./web
         run: npx playwright test

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

@@ -48,6 +48,8 @@ jobs:
           -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"
             
     - 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_IS_CLOUD: ${{ secrets.CONFLUENCE_IS_CLOUD }}
   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 }}
 

.gitignore

@@ -17,12 +17,24 @@ 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,6 +23,9 @@ 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
 
@@ -46,7 +49,6 @@ PYTHONUNBUFFERED=1
 
 
 # Internet Search
-BING_API_KEY=<REPLACE THIS>
 EXA_API_KEY=<REPLACE THIS>
 
 
@@ -65,3 +67,12 @@ 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,21 +24,23 @@
           "Celery primary",
           "Celery light",
           "Celery heavy",
-          "Celery indexing",
-          "Celery user files indexing",
+          "Celery docfetching",
+          "Celery docprocessing",
           "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)",
@@ -46,14 +48,15 @@
           "Celery primary",
           "Celery light",
           "Celery heavy",
-          "Celery indexing",
-          "Celery user files indexing",
+          "Celery docfetching",
+          "Celery docprocessing",
           "Celery beat",
           "Celery monitoring"
         ],
         "presentation": {
           "group": "1"
-        }
+        },
+        "stopAll": true
       }
     ],
     "configurations": [
@@ -189,7 +192,7 @@
           "--loglevel=INFO",
           "--hostname=light@%n",
           "-Q",
-          "vespa_metadata_sync,connector_deletion,doc_permissions_upsert"
+          "vespa_metadata_sync,connector_deletion,doc_permissions_upsert,index_attempt_cleanup"
         ],
         "presentation": {
           "group": "2"
@@ -226,35 +229,66 @@
         "consoleTitle": "Celery heavy Console"
       },
       {
-        "name": "Celery indexing",
+        "name": "Celery docfetching",
         "type": "debugpy",
         "request": "launch",
         "module": "celery",
         "cwd": "${workspaceFolder}/backend",
         "envFile": "${workspaceFolder}/.vscode/.env",
         "env": {
-          "ENABLE_MULTIPASS_INDEXING": "false",
-          "LOG_LEVEL": "DEBUG",
-          "PYTHONUNBUFFERED": "1",
-          "PYTHONPATH": "."
+            "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=indexing@%n",
-          "-Q",
-          "connector_indexing"
+            "-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"
         ],
         "presentation": {
-          "group": "2"
+            "group": "2"
         },
-        "consoleTitle": "Celery indexing Console"
-      },
+        "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
+    },
       {
         "name": "Celery monitoring",
         "type": "debugpy",
@@ -303,35 +337,6 @@
         },
         "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",
@@ -426,7 +431,7 @@
       },
       "args": [
         "--filename",
-        "generated/openapi.json",
+        "generated/openapi.json"
       ]
     },
     {

CONTRIBUTING.md

@@ -103,10 +103,10 @@ If using PowerShell, the command slightly differs:
 Install the required python dependencies:
 
 ```bash
-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
+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
 ```
 
 Install Playwright for Python (headless browser required by the Web Connector)

CONTRIBUTING_VSCODE.md

@@ -5,7 +5,7 @@ This guide explains how to set up and use VSCode's debugging capabilities with t
 ## Initial Setup
 
 1. **Environment Setup**:
-   - Copy `.vscode/.env.template` to `.vscode/.env`
+   - Copy `.vscode/env_template.txt` to `.vscode/.env`
    - Fill in the necessary environment variables in `.vscode/.env`
 2. **launch.json**:
    - Copy `.vscode/launch.template.jsonc` to `.vscode/launch.json`
@@ -17,10 +17,9 @@ Before starting, make sure the Docker Daemon is running.
 1. Open the Debug view in VSCode (Cmd+Shift+D on macOS)
 2. From the dropdown at the top, select "Clear and Restart External Volumes and Containers" and press the green play button
 3. From the dropdown at the top, select "Run All Onyx Services" and press the green play button
-4. CD into web, run "npm i" followed by npm run dev.
-5. Now, you can navigate to onyx in your browser (default is http://localhost:3000) and start using the app
-6. You can set breakpoints by clicking to the left of line numbers to help debug while the app is running
-7. Use the debug toolbar to step through code, inspect variables, etc.
+4. Now, you can navigate to onyx in your browser (default is http://localhost:3000) and start using the app
+5. You can set breakpoints by clicking to the left of line numbers to help debug while the app is running
+6. Use the debug toolbar to step through code, inspect variables, etc.
 
 ## Features
 

backend/Dockerfile

@@ -12,7 +12,8 @@ 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"
+    DO_NOT_TRACK="true" \
+    PLAYWRIGHT_BROWSERS_PATH="/app/.cache/ms-playwright"
 
 
 RUN echo "ONYX_VERSION: ${ONYX_VERSION}"
@@ -116,6 +117,14 @@ 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,11 +9,20 @@ 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"
-
+    DANSWER_RUNNING_IN_DOCKER="true" \
+    HF_HOME=/app/.cache/huggingface
 
 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 \
@@ -38,9 +47,11 @@ 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 /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
+# 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
 
 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_STANDARD_VALUE,
+    POSTGRES_DEFAULT_SCHEMA,
     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_STANDARD_VALUE]
+        schemas = [POSTGRES_DEFAULT_SCHEMA]
 
     # without init_engine, subsequent engine calls fail hard intentionally
     SqlEngine.init_engine(pool_size=20, max_overflow=5)

backend/alembic/versions/12635f6655b7_drive_canonical_ids.py

@@ -96,7 +96,7 @@ def get_google_drive_documents_from_database() -> list[dict]:
     result = bind.execute(
         sa.text(
             """
-            SELECT d.id, cc.id as cc_pair_id
+            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
@@ -109,7 +109,7 @@ def get_google_drive_documents_from_database() -> list[dict]:
 
     documents = []
     for row in result:
-        documents.append({"document_id": row.id, "cc_pair_id": row.cc_pair_id})
+        documents.append({"document_id": row.id})
 
     return documents
 

backend/alembic/versions/2f95e36923e6_add_indexing_coordination.py

@@ -0,0 +1,115 @@
+"""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_STANDARD_VALUE
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 
 # 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_STANDARD_VALUE}.show_trgm(cleaned_name);
+                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.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_STANDARD_VALUE}.show_trgm(cleaned_name)
+                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name)
                 WHERE document_id = NEW.id;
                 RETURN NEW;
             END;

backend/alembic/versions/3fc5d75723b3_add_doc_metadata_field_in_document_model.py

@@ -0,0 +1,30 @@
+"""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_STANDARD_VALUE
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 
 
 # revision identifiers, used by Alembic.
@@ -478,7 +478,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_STANDARD_VALUE}.gin_trgm_ops)"
+        f"ON kg_entity USING GIN (name {POSTGRES_DEFAULT_SCHEMA}.gin_trgm_ops)"
     )
     op.execute(
         "CREATE INDEX IF NOT EXISTS idx_kg_entity_normalization_trigrams "
@@ -518,7 +518,7 @@ def upgrade() -> None:
 
                 -- Set name and name trigrams
                 NEW.name = name;
-                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.show_trgm(cleaned_name);
+                NEW.name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name);
                 RETURN NEW;
             END;
             $$ LANGUAGE plpgsql;
@@ -563,7 +563,7 @@ def upgrade() -> None:
                 UPDATE kg_entity
                 SET
                     name = doc_name,
-                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE}.show_trgm(cleaned_name)
+                    name_trigrams = {POSTGRES_DEFAULT_SCHEMA}.show_trgm(cleaned_name)
                 WHERE document_id = NEW.id;
                 RETURN NEW;
             END;

backend/alembic/versions/5ae8240accb3_add_research_agent_database_tables_and_.py

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

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

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

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

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

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

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

@@ -0,0 +1,147 @@
+"""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(db_session=session)
+    external_store = get_s3_file_store()
 
     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(db_session=session)
+    external_store = get_s3_file_store()
 
     # Find all files currently stored in PostgreSQL (lobj_oid is not null)
     result = session.execute(

backend/alembic/versions/f8a9b2c3d4e5_add_research_answer_purpose_to_chat_message.py

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

@@ -0,0 +1,69 @@
+"""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(db_session).save_file(
+            get_default_file_store().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,6 +47,7 @@ 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
@@ -58,7 +59,9 @@ 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
@@ -422,7 +425,7 @@ def connector_permission_sync_generator_task(
 
     lock: RedisLock = r.lock(
         OnyxRedisLocks.CONNECTOR_DOC_PERMISSIONS_SYNC_LOCK_PREFIX
-        + f"_{redis_connector.id}",
+        + f"_{redis_connector.cc_pair_id}",
         timeout=CELERY_PERMISSIONS_SYNC_LOCK_TIMEOUT,
         thread_local=False,
     )
@@ -498,16 +501,31 @@ 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() -> list[str]:
-                return get_document_ids_for_connector_credential_pair(
+            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(
                     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, callback
+                cc_pair,
+                fetch_all_existing_docs_fn,
+                fetch_all_existing_docs_ids_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.id}",
+        + f"_{redis_connector.cc_pair_id}",
         timeout=CELERY_EXTERNAL_GROUP_SYNC_LOCK_TIMEOUT,
     )
 

backend/ee/onyx/chat/process_message.py

@@ -1,38 +0,0 @@
-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,6 +71,19 @@ 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
 #####
@@ -89,6 +102,19 @@ 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,7 +114,6 @@ def get_all_usage_reports(db_session: Session) -> list[UsageReportMetadata]:
 
 
 def get_usage_report_data(
-    db_session: Session,
     report_display_name: str,
 ) -> IO:
     """
@@ -128,7 +127,7 @@ def get_usage_report_data(
     Returns:
         The usage report data.
     """
-    file_store = get_default_file_store(db_session)
+    file_store = get_default_file_store()
     # 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/document_index/vespa/app_config/cloud-services.xml.jinja

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

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

@@ -6,6 +6,7 @@ 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
@@ -25,6 +26,7 @@ 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]:
     """
@@ -43,7 +45,7 @@ def confluence_doc_sync(
 
     yield from generic_doc_sync(
         cc_pair=cc_pair,
-        fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,
+        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,
         callback=callback,
         doc_source=DocumentSource.CONFLUENCE,
         slim_connector=confluence_connector,

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

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

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

@@ -0,0 +1,488 @@
+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,6 +3,7 @@ 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
@@ -35,6 +36,7 @@ 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,6 +8,7 @@ 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
@@ -169,6 +170,7 @@ 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/jira/doc_sync.py

@@ -1,6 +1,7 @@
 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
@@ -17,6 +18,7 @@ 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(
@@ -26,7 +28,7 @@ def jira_doc_sync(
 
     yield from generic_doc_sync(
         cc_pair=cc_pair,
-        fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,
+        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_fn,
         callback=callback,
         doc_source=DocumentSource.JIRA,
         slim_connector=jira_connector,

backend/ee/onyx/external_permissions/perm_sync_types.py

@@ -2,27 +2,45 @@ 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
-
-# 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
+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
 
 
 class FetchAllDocumentsFunction(Protocol):
-    """Protocol for a function that fetches all document IDs for a connector credential pair."""
+    """Protocol for a function that fetches documents for a connector credential pair.
 
-    def __call__(self) -> list[str]:
+    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]:
         """
-        Returns a list of document IDs for a connector credential pair.
+        Fetches documents for a connector credential pair.
+        """
+        ...
 
-        This is typically used to determine which documents should no longer be
-        accessible during the document sync process.
+
+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.
         """
         ...
 
@@ -30,19 +48,20 @@ class FetchAllDocumentsFunction(Protocol):
 # Defining the input/output types for the sync functions
 DocSyncFuncType = Callable[
     [
-        "ConnectorCredentialPair",
+        ConnectorCredentialPair,
         FetchAllDocumentsFunction,
-        Optional["IndexingHeartbeatInterface"],
+        FetchAllDocumentsIdsFunction,
+        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

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

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

@@ -0,0 +1,684 @@
+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,6 +3,7 @@ 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
@@ -130,6 +131,7 @@ 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,12 +7,18 @@ 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
@@ -20,10 +26,13 @@ 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
@@ -63,6 +72,7 @@ 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"""
@@ -117,6 +127,18 @@ _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,
@@ -138,6 +160,18 @@ _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,6 +1,7 @@
 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,6 +19,7 @@ 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(
@@ -27,7 +29,7 @@ def teams_doc_sync(
 
     yield from generic_doc_sync(
         cc_pair=cc_pair,
-        fetch_all_existing_docs_fn=fetch_all_existing_docs_fn,
+        fetch_all_existing_docs_ids_fn=fetch_all_existing_docs_ids_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 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.configs.constants import DocumentSource
@@ -14,7 +14,7 @@ logger = setup_logger()
 
 def generic_doc_sync(
     cc_pair: ConnectorCredentialPair,
-    fetch_all_existing_docs_fn: FetchAllDocumentsFunction,
+    fetch_all_existing_docs_ids_fn: FetchAllDocumentsIdsFunction,
     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 = set(fetch_all_existing_docs_fn())
+    existing_doc_ids: list[str] = fetch_all_existing_docs_ids_fn()
 
-    missing_doc_ids = existing_doc_ids - newly_fetched_doc_ids
+    missing_doc_ids = set(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_bot_msg=message_info.is_bot_msg,
+            is_slash_command=message_info.is_slash_command,
         )
 
         answer_blocks = build_standard_answer_blocks(

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

@@ -134,15 +134,14 @@ 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, db_session=db_session, is_logotype=is_logotype)
+    upload_logo(file=file, is_logotype=is_logotype)
 
 
 def fetch_logo_helper(db_session: Session) -> Response:
     try:
-        file_store = get_default_file_store(db_session)
+        file_store = get_default_file_store()
         onyx_file = file_store.get_file_with_mime_type(get_logo_filename())
         if not onyx_file:
             raise ValueError("get_onyx_file returned None!")
@@ -158,7 +157,7 @@ def fetch_logo_helper(db_session: Session) -> Response:
 
 def fetch_logotype_helper(db_session: Session) -> Response:
     try:
-        file_store = get_default_file_store(db_session)
+        file_store = get_default_file_store()
         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,7 +6,6 @@ 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
@@ -99,9 +98,7 @@ def guess_file_type(filename: str) -> str:
     return "application/octet-stream"
 
 
-def upload_logo(
-    db_session: Session, file: UploadFile | str, is_logotype: bool = False
-) -> bool:
+def upload_logo(file: UploadFile | str, is_logotype: bool = False) -> bool:
     content: IO[Any]
 
     if isinstance(file, str):
@@ -129,7 +126,7 @@ def upload_logo(
         display_name = file.filename
         file_type = file.content_type or "image/jpeg"
 
-    file_store = get_default_file_store(db_session)
+    file_store = get_default_file_store()
     file_store.save_file(
         content=content,
         display_name=display_name,

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

@@ -1,42 +1,22 @@
-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 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.models import ChatBasicResponse
+from onyx.chat.process_message import gather_stream
 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
@@ -45,7 +25,6 @@ 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
 
@@ -54,177 +33,6 @@ 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,
@@ -237,13 +45,36 @@ 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_message_req.chat_session_id, db_session=db_session
+            chat_session_id=chat_session_id, db_session=db_session
         )
     except Exception:
         parent_message = get_or_create_root_message(
-            chat_session_id=chat_message_req.chat_session_id, db_session=db_session
+            chat_session_id=chat_session_id, db_session=db_session
         )
 
     if (
@@ -258,7 +89,7 @@ def handle_simplified_chat_message(
         retrieval_options = chat_message_req.retrieval_options
 
     full_chat_msg_info = CreateChatMessageRequest(
-        chat_session_id=chat_message_req.chat_session_id,
+        chat_session_id=chat_session_id,
         parent_message_id=parent_message.id,
         message=chat_message_req.message,
         file_descriptors=[],
@@ -283,7 +114,7 @@ def handle_simplified_chat_message(
         enforce_chat_session_id_for_search_docs=False,
     )
 
-    return _convert_packet_stream_to_response(packets)
+    return gather_stream(packets)
 
 
 @router.post("/send-message-simple-with-history")
@@ -403,4 +234,4 @@ def handle_send_message_simple_with_history(
         enforce_chat_session_id_for_search_docs=False,
     )
 
-    return _convert_packet_stream_to_response(packets)
+    return gather_stream(packets)

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

@@ -6,10 +6,8 @@ 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
@@ -17,8 +15,9 @@ 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):
@@ -41,11 +40,13 @@ class DocumentSearchRequest(ChunkContext):
 
 
 class BasicCreateChatMessageRequest(ChunkContext):
-    """Before creating messages, be sure to create a chat_session and get an id
+    """If a chat_session_id is not provided, a persona_id must be provided to automatically create a new chat session
     Note, for simplicity this option only allows for a single linear chain of messages
     """
 
-    chat_session_id: UUID
+    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
     # New message contents
     message: str
     # Defaults to using retrieval with no additional filters
@@ -62,6 +63,12 @@ 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
@@ -148,30 +155,6 @@ 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
@@ -182,7 +165,6 @@ 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
@@ -214,6 +196,5 @@ 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,7 +8,6 @@ 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,
 )
@@ -20,8 +19,10 @@ 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.process_message import ChatPacketStream
+from onyx.chat.models import QADocsResponse
+from onyx.chat.process_message import gather_stream
 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
@@ -39,6 +40,7 @@ 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
 
@@ -140,7 +142,7 @@ def get_answer_stream(
     query_request: OneShotQARequest,
     user: User | None = Depends(current_user),
     db_session: Session = Depends(get_session),
-) -> ChatPacketStream:
+) -> AnswerStream:
     query = query_request.messages[0].message
     logger.notice(f"Received query for Answer API: {query}")
 
@@ -205,7 +207,6 @@ def get_answer_stream(
         new_msg_req=request,
         user=user,
         db_session=db_session,
-        include_contexts=query_request.return_contexts,
     )
 
     return packets
@@ -219,12 +220,28 @@ def get_answer_with_citation(
 ) -> OneShotQAResponse:
     try:
         packets = get_answer_stream(request, user, db_session)
-        answer = gather_stream_for_answer_api(packets)
+        answer = gather_stream(packets)
 
         if answer.error_msg:
             raise RuntimeError(answer.error_msg)
 
-        return answer
+        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,
+            ),
+        )
     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(db_session)
+    file_store = get_default_file_store()
 
     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(db_session)
+    file_store = get_default_file_store()
     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(db_session, report_name)
+        file = get_usage_report_data(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.OTHER,
+            file_origin=FileOrigin.GENERATED_REPORT,
             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.OTHER,
+            file_origin=FileOrigin.GENERATED_REPORT,
             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(db_session)
+    file_store = get_default_file_store()
 
     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(db_session: Session, logo_path: str | None) -> None:
+def _seed_logo(logo_path: str | None) -> None:
     if logo_path:
         logger.notice("Uploading logo")
-        upload_logo(db_session=db_session, file=logo_path)
+        upload_logo(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(db_session, seed_config.seeded_logo_path)
+        _seed_logo(seed_config.seeded_logo_path)
         _seed_enterprise_settings(seed_config)
         _seed_analytics_script(seed_config)
 

backend/model_server/constants.py

@@ -1,34 +1,5 @@
-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,55 +1,30 @@
 import asyncio
-import json
 import time
-from types import TracebackType
-from typing import cast
+from typing import Any
 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
 
@@ -57,315 +32,6 @@ _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,
@@ -404,20 +70,34 @@ 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):
@@ -434,52 +114,10 @@ async def embed_text(
     for text in texts:
         total_chars += len(text)
 
-    if provider_type is not None:
-        logger.info(
-            f"Embedding {len(texts)} texts with {total_chars} total characters with provider: {provider_type}"
-        )
+    # Only local models should call this function now
+    # API providers should go directly to API server
 
-        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:
+    if model_name is not None:
         logger.info(
             f"Embedding {len(texts)} texts with {total_chars} total characters with local model: {model_name}"
         )
@@ -492,8 +130,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: local_model.encode(
-                prefixed_texts, normalize_embeddings=normalize_embeddings
+            lambda: _concurrent_embedding(
+                prefixed_texts, local_model, normalize_embeddings
             ),
         )
         embeddings = [
@@ -515,10 +153,8 @@ async def embed_text(
             f"elapsed={elapsed:.2f}"
         )
     else:
-        logger.error("Neither model name nor provider specified for embedding")
-        raise ValueError(
-            "Either model name or provider must be provided to run embeddings."
-        )
+        logger.error("Model name not specified for embedding")
+        raise ValueError("Model name must be provided to run embeddings.")
 
     return embeddings
 
@@ -533,77 +169,6 @@ 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,
@@ -615,6 +180,13 @@ 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")
 
@@ -632,26 +204,12 @@ 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,
@@ -669,6 +227,13 @@ 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")
 
@@ -680,55 +245,13 @@ async def process_rerank_request(rerank_request: RerankRequest) -> RerankRespons
         raise ValueError("Empty documents cannot be reranked.")
 
     try:
-        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}")
+        # 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)
 
     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(os.path.expanduser("~")) / ".cache/huggingface"
-TEMP_HF_CACHE_PATH = Path(os.path.expanduser("~")) / ".cache/temp_huggingface"
+HF_CACHE_PATH = Path(".cache/huggingface")
+TEMP_HF_CACHE_PATH = Path(".cache/temp_huggingface")
 
 transformer_logging.set_verbosity_error()
 

backend/model_server/utils.py

@@ -70,32 +70,3 @@ 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

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

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

@@ -1,64 +0,0 @@
-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,6 +10,7 @@ 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,8 +14,6 @@ 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
@@ -139,17 +137,6 @@ 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,8 +9,6 @@ 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()
@@ -23,17 +21,6 @@ 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,8 +12,6 @@ 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
@@ -33,17 +31,6 @@ 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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,81 +0,0 @@
-from langgraph.graph import END
-from langgraph.graph import START
-from langgraph.graph import StateGraph
-
-from onyx.agents.agent_search.deep_search.initial.retrieve_orig_question_docs.nodes.format_orig_question_search_input import (
-    format_orig_question_search_input,
-)
-from onyx.agents.agent_search.deep_search.initial.retrieve_orig_question_docs.nodes.format_orig_question_search_output import (
-    format_orig_question_search_output,
-)
-from onyx.agents.agent_search.deep_search.initial.retrieve_orig_question_docs.states import (
-    BaseRawSearchInput,
-)
-from onyx.agents.agent_search.deep_search.initial.retrieve_orig_question_docs.states import (
-    BaseRawSearchOutput,
-)
-from onyx.agents.agent_search.deep_search.initial.retrieve_orig_question_docs.states import (
-    BaseRawSearchState,
-)
-from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.graph_builder import (
-    expanded_retrieval_graph_builder,
-)
-
-
-def retrieve_orig_question_docs_graph_builder() -> StateGraph:
-    """
-    LangGraph graph builder for the retrieval of documents
-    that are relevant to the original question. This is
-    largely a wrapper around the expanded retrieval process to
-    ensure parallelism with the sub-question answer process.
-    """
-    graph = StateGraph(
-        state_schema=BaseRawSearchState,
-        input=BaseRawSearchInput,
-        output=BaseRawSearchOutput,
-    )
-
-    ### Add nodes ###
-
-    # Format the original question search output
-    graph.add_node(
-        node="format_orig_question_search_output",
-        action=format_orig_question_search_output,
-    )
-
-    # The sub-graph that executes the expanded retrieval process
-    expanded_retrieval = expanded_retrieval_graph_builder().compile()
-    graph.add_node(
-        node="retrieve_orig_question_docs_subgraph",
-        action=expanded_retrieval,
-    )
-
-    # Format the original question search input
-    graph.add_node(
-        node="format_orig_question_search_input",
-        action=format_orig_question_search_input,
-    )
-
-    ### Add edges ###
-
-    graph.add_edge(start_key=START, end_key="format_orig_question_search_input")
-
-    graph.add_edge(
-        start_key="format_orig_question_search_input",
-        end_key="retrieve_orig_question_docs_subgraph",
-    )
-    graph.add_edge(
-        start_key="retrieve_orig_question_docs_subgraph",
-        end_key="format_orig_question_search_output",
-    )
-
-    graph.add_edge(
-        start_key="format_orig_question_search_output",
-        end_key=END,
-    )
-
-    return graph
-
-
-if __name__ == "__main__":
-    pass

backend/onyx/agents/agent_search/deep_search/initial/retrieve_orig_question_docs/nodes/format_orig_question_search_input.py

@@ -1,28 +0,0 @@
-from typing import cast
-
-from langchain_core.runnables.config import RunnableConfig
-
-from onyx.agents.agent_search.core_state import CoreState
-from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states import (
-    ExpandedRetrievalInput,
-)
-from onyx.agents.agent_search.models import GraphConfig
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def format_orig_question_search_input(
-    state: CoreState, config: RunnableConfig
-) -> ExpandedRetrievalInput:
-    """
-    LangGraph node to format the search input for the original question.
-    """
-    logger.debug("generate_raw_search_data")
-    graph_config = cast(GraphConfig, config["metadata"]["config"])
-    return ExpandedRetrievalInput(
-        question=graph_config.inputs.prompt_builder.raw_user_query,
-        base_search=True,
-        sub_question_id=None,  # This graph is always and only used for the original question
-        log_messages=[],
-    )

backend/onyx/agents/agent_search/deep_search/initial/retrieve_orig_question_docs/nodes/format_orig_question_search_output.py

@@ -1,30 +0,0 @@
-from onyx.agents.agent_search.deep_search.main.states import OrigQuestionRetrievalUpdate
-from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states import (
-    ExpandedRetrievalOutput,
-)
-from onyx.agents.agent_search.shared_graph_utils.models import AgentChunkRetrievalStats
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def format_orig_question_search_output(
-    state: ExpandedRetrievalOutput,
-) -> OrigQuestionRetrievalUpdate:
-    """
-    LangGraph node to format the search result for the original question into the
-    proper format.
-    """
-    sub_question_retrieval_stats = state.expanded_retrieval_result.retrieval_stats
-    if sub_question_retrieval_stats is None:
-        sub_question_retrieval_stats = AgentChunkRetrievalStats()
-    else:
-        sub_question_retrieval_stats = sub_question_retrieval_stats
-
-    return OrigQuestionRetrievalUpdate(
-        orig_question_verified_reranked_documents=state.expanded_retrieval_result.verified_reranked_documents,
-        orig_question_sub_query_retrieval_results=state.expanded_retrieval_result.expanded_query_results,
-        orig_question_retrieved_documents=state.retrieved_documents,
-        orig_question_retrieval_stats=sub_question_retrieval_stats,
-        log_messages=[],
-    )

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

@@ -1,29 +0,0 @@
-from onyx.agents.agent_search.deep_search.main.states import (
-    OrigQuestionRetrievalUpdate,
-)
-from onyx.agents.agent_search.deep_search.shared.expanded_retrieval.states import (
-    ExpandedRetrievalInput,
-)
-
-
-## Graph Input State
-class BaseRawSearchInput(ExpandedRetrievalInput):
-    pass
-
-
-## Graph Output State
-class BaseRawSearchOutput(OrigQuestionRetrievalUpdate):
-    """
-    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.
-    """
-
-    # base_expanded_retrieval_result: QuestionRetrievalResult = QuestionRetrievalResult()
-
-
-## Graph State
-class BaseRawSearchState(
-    BaseRawSearchInput, BaseRawSearchOutput, OrigQuestionRetrievalUpdate
-):
-    pass

backend/onyx/agents/agent_search/deep_search/main/edges.py

@@ -1,113 +0,0 @@
-from collections.abc import Hashable
-from datetime import datetime
-from typing import cast
-from typing import Literal
-
-from langchain_core.runnables import RunnableConfig
-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.main.states import MainState
-from onyx.agents.agent_search.deep_search.main.states import (
-    RequireRefinemenEvalUpdate,
-)
-from onyx.agents.agent_search.models import GraphConfig
-from onyx.agents.agent_search.shared_graph_utils.utils import make_question_id
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-
-def route_initial_tool_choice(
-    state: MainState, config: RunnableConfig
-) -> Literal["call_tool", "start_agent_search", "logging_node"]:
-    """
-    LangGraph edge to route to agent search.
-    """
-    agent_config = cast(GraphConfig, config["metadata"]["config"])
-    if state.tool_choice is None:
-        return "logging_node"
-
-    if (
-        agent_config.behavior.use_agentic_search
-        and agent_config.tooling.search_tool is not None
-        and state.tool_choice.tool.name == agent_config.tooling.search_tool.name
-    ):
-        return "start_agent_search"
-    else:
-        return "call_tool"
-
-
-def parallelize_initial_sub_question_answering(
-    state: MainState,
-) -> list[Send | Hashable]:
-    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=[],
-                ),
-            )
-        ]
-
-
-# Define the function that determines whether to continue or not
-def continue_to_refined_answer_or_end(
-    state: RequireRefinemenEvalUpdate,
-) -> Literal["create_refined_sub_questions", "logging_node"]:
-    if state.require_refined_answer_eval:
-        return "create_refined_sub_questions"
-    else:
-        return "logging_node"
-
-
-def parallelize_refined_sub_question_answering(
-    state: MainState,
-) -> list[Send | Hashable]:
-    edge_start_time = datetime.now()
-    if len(state.refined_sub_questions) > 0:
-        return [
-            Send(
-                "answer_refined_question_subgraphs",
-                SubQuestionAnsweringInput(
-                    question=question_data.sub_question,
-                    question_id=make_question_id(1, question_num),
-                    log_messages=[
-                        f"{edge_start_time} -- Main Edge - Parallelize Refined Sub-question Answering"
-                    ],
-                ),
-            )
-            for question_num, question_data in state.refined_sub_questions.items()
-        ]
-
-    else:
-        return [
-            Send(
-                "ingest_refined_sub_answers",
-                AnswerQuestionOutput(
-                    answer_results=[],
-                ),
-            )
-        ]

backend/onyx/agents/agent_search/deep_search/main/graph_builder.py

@@ -1,263 +0,0 @@
-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.graph_builder import (
-    generate_initial_answer_graph_builder,
-)
-from onyx.agents.agent_search.deep_search.main.edges import (
-    continue_to_refined_answer_or_end,
-)
-from onyx.agents.agent_search.deep_search.main.edges import (
-    parallelize_refined_sub_question_answering,
-)
-from onyx.agents.agent_search.deep_search.main.edges import (
-    route_initial_tool_choice,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.compare_answers import (
-    compare_answers,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.create_refined_sub_questions import (
-    create_refined_sub_questions,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.decide_refinement_need import (
-    decide_refinement_need,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.extract_entities_terms import (
-    extract_entities_terms,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.generate_validate_refined_answer import (
-    generate_validate_refined_answer,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.ingest_refined_sub_answers import (
-    ingest_refined_sub_answers,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.persist_agent_results import (
-    persist_agent_results,
-)
-from onyx.agents.agent_search.deep_search.main.nodes.start_agent_search import (
-    start_agent_search,
-)
-from onyx.agents.agent_search.deep_search.main.states import MainInput
-from onyx.agents.agent_search.deep_search.main.states import MainState
-from onyx.agents.agent_search.deep_search.refinement.consolidate_sub_answers.graph_builder import (
-    answer_refined_query_graph_builder,
-)
-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.agents.agent_search.shared_graph_utils.utils import get_test_config
-from onyx.utils.logger import setup_logger
-
-logger = setup_logger()
-
-test_mode = False
-
-
-def agent_search_graph_builder() -> StateGraph:
-    """
-    LangGraph graph builder for the main agent search process.
-    """
-    graph = StateGraph(
-        state_schema=MainState,
-        input=MainInput,
-    )
-
-    # Prepare the tool input
-    graph.add_node(
-        node="prepare_tool_input",
-        action=prepare_tool_input,
-    )
-
-    # Choose the initial tool
-    graph.add_node(
-        node="choose_tool",
-        action=choose_tool,
-    )
-
-    # Call the tool, if required
-    graph.add_node(
-        node="call_tool",
-        action=call_tool,
-    )
-
-    # Use the tool response
-    graph.add_node(
-        node="basic_use_tool_response",
-        action=basic_use_tool_response,
-    )
-
-    # Start the agent search process
-    graph.add_node(
-        node="start_agent_search",
-        action=start_agent_search,
-    )
-
-    # The sub-graph for the initial answer generation
-    generate_initial_answer_subgraph = generate_initial_answer_graph_builder().compile()
-    graph.add_node(
-        node="generate_initial_answer_subgraph",
-        action=generate_initial_answer_subgraph,
-    )
-
-    # Create the refined sub-questions
-    graph.add_node(
-        node="create_refined_sub_questions",
-        action=create_refined_sub_questions,
-    )
-
-    # Subgraph for the refined sub-answer generation
-    answer_refined_question = answer_refined_query_graph_builder().compile()
-    graph.add_node(
-        node="answer_refined_question_subgraphs",
-        action=answer_refined_question,
-    )
-
-    # Ingest the refined sub-answers
-    graph.add_node(
-        node="ingest_refined_sub_answers",
-        action=ingest_refined_sub_answers,
-    )
-
-    # Node to generate the refined answer
-    graph.add_node(
-        node="generate_validate_refined_answer",
-        action=generate_validate_refined_answer,
-    )
-
-    # Early node to extract the entities and terms from the initial answer,
-    # This information is used to inform the creation the refined sub-questions
-    graph.add_node(
-        node="extract_entity_term",
-        action=extract_entities_terms,
-    )
-
-    # Decide if the answer needs to be refined (currently always true)
-    graph.add_node(
-        node="decide_refinement_need",
-        action=decide_refinement_need,
-    )
-
-    # Compare the initial and refined answers, and determine whether
-    # the refined answer is sufficiently better
-    graph.add_node(
-        node="compare_answers",
-        action=compare_answers,
-    )
-
-    # Log the results. This will log the stats as well as the answers, sub-questions, and sub-answers
-    graph.add_node(
-        node="logging_node",
-        action=persist_agent_results,
-    )
-
-    ### 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",
-        route_initial_tool_choice,
-        ["call_tool", "start_agent_search", "logging_node"],
-    )
-
-    graph.add_edge(
-        start_key="call_tool",
-        end_key="basic_use_tool_response",
-    )
-    graph.add_edge(
-        start_key="basic_use_tool_response",
-        end_key="logging_node",
-    )
-
-    graph.add_edge(
-        start_key="start_agent_search",
-        end_key="generate_initial_answer_subgraph",
-    )
-
-    graph.add_edge(
-        start_key="start_agent_search",
-        end_key="extract_entity_term",
-    )
-
-    # Wait for the initial answer generation and the entity/term extraction to be complete
-    # before deciding if a refinement is needed.
-    graph.add_edge(
-        start_key=["generate_initial_answer_subgraph", "extract_entity_term"],
-        end_key="decide_refinement_need",
-    )
-
-    graph.add_conditional_edges(
-        source="decide_refinement_need",
-        path=continue_to_refined_answer_or_end,
-        path_map=["create_refined_sub_questions", "logging_node"],
-    )
-
-    graph.add_conditional_edges(
-        source="create_refined_sub_questions",
-        path=parallelize_refined_sub_question_answering,
-        path_map=["answer_refined_question_subgraphs"],
-    )
-    graph.add_edge(
-        start_key="answer_refined_question_subgraphs",
-        end_key="ingest_refined_sub_answers",
-    )
-
-    graph.add_edge(
-        start_key="ingest_refined_sub_answers",
-        end_key="generate_validate_refined_answer",
-    )
-
-    graph.add_edge(
-        start_key="generate_validate_refined_answer",
-        end_key="compare_answers",
-    )
-    graph.add_edge(
-        start_key="compare_answers",
-        end_key="logging_node",
-    )
-
-    graph.add_edge(
-        start_key="logging_node",
-        end_key=END,
-    )
-
-    return graph
-
-
-if __name__ == "__main__":
-    pass
-
-    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 = agent_search_graph_builder()
-    compiled_graph = graph.compile()
-    primary_llm, fast_llm = get_default_llms()
-
-    with get_session_with_current_tenant() as db_session:
-        search_request = SearchRequest(query="Who created Excel?")
-        graph_config = get_test_config(
-            db_session, primary_llm, fast_llm, search_request
-        )
-
-        inputs = MainInput(log_messages=[])
-
-        for thing in compiled_graph.stream(
-            input=inputs,
-            config={"configurable": {"config": graph_config}},
-            stream_mode="custom",
-            subgraphs=True,
-        ):
-            logger.debug(thing)