Merge pull request #3161 from danswer-ai/hotfix/v0.13-indexing-redux

enhanced logging for indexing and increased indexing timeouts

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

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

backend/alembic/versions/b156fa702355_chat_reworked.py

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

backend/alembic/versions/c99d76fcd298_add_nullable_to_persona_id_in_chat_.py

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

backend/danswer/auth/users.py

@@ -100,6 +100,11 @@ from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 logger = setup_logger()
 
 
+class BasicAuthenticationError(HTTPException):
+    def __init__(self, detail: str):
+        super().__init__(status_code=status.HTTP_403_FORBIDDEN, detail=detail)
+
+
 def is_user_admin(user: User | None) -> bool:
     if AUTH_TYPE == AuthType.DISABLED:
         return True
@@ -463,8 +468,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             has_web_login = attributes.get_attribute(user, "has_web_login")
 
             if not has_web_login:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
+                raise BasicAuthenticationError(
                     detail="NO_WEB_LOGIN_AND_HAS_NO_PASSWORD",
                 )
 
@@ -621,14 +625,12 @@ async def double_check_user(
         return None
 
     if user is None:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not authenticated.",
         )
 
     if user_needs_to_be_verified() and not user.is_verified:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not verified.",
         )
 
@@ -637,8 +639,7 @@ async def double_check_user(
         and user.oidc_expiry < datetime.now(timezone.utc)
         and not include_expired
     ):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User's OIDC token has expired.",
         )
 
@@ -664,15 +665,13 @@ async def current_curator_or_admin_user(
         return None
 
     if not user or not hasattr(user, "role"):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not authenticated or lacks role information.",
         )
 
     allowed_roles = {UserRole.GLOBAL_CURATOR, UserRole.CURATOR, UserRole.ADMIN}
     if user.role not in allowed_roles:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User is not a curator or admin.",
         )
 
@@ -684,8 +683,7 @@ async def current_admin_user(user: User | None = Depends(current_user)) -> User
         return None
 
     if not user or not hasattr(user, "role") or user.role != UserRole.ADMIN:
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User must be an admin to perform this action.",
         )
 

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

@@ -59,7 +59,7 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
     logger.info(f"Multiprocessing start method: {multiprocessing.get_start_method()}")
 
     SqlEngine.set_app_name(POSTGRES_CELERY_WORKER_INDEXING_APP_NAME)
-    SqlEngine.init_engine(pool_size=8, max_overflow=0)
+    SqlEngine.init_engine(pool_size=sender.concurrency, max_overflow=sender.concurrency)
 
     # Startup checks are not needed in multi-tenant case
     if MULTI_TENANT:

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

@@ -14,10 +14,14 @@ from celery.signals import worker_shutdown
 import danswer.background.celery.apps.app_base as app_base
 from danswer.background.celery.apps.app_base import task_logger
 from danswer.background.celery.celery_utils import celery_is_worker_primary
+from danswer.background.celery.tasks.vespa.tasks import get_unfenced_index_attempt_ids
 from danswer.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT
 from danswer.configs.constants import DanswerRedisLocks
 from danswer.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME
+from danswer.db.engine import get_session_with_default_tenant
 from danswer.db.engine import SqlEngine
+from danswer.db.index_attempt import get_index_attempt
+from danswer.db.index_attempt import mark_attempt_failed
 from danswer.redis.redis_connector_credential_pair import RedisConnectorCredentialPair
 from danswer.redis.redis_connector_delete import RedisConnectorDelete
 from danswer.redis.redis_connector_index import RedisConnectorIndex
@@ -134,6 +138,23 @@ def on_worker_init(sender: Any, **kwargs: Any) -> None:
 
     RedisConnectorStop.reset_all(r)
 
+    # mark orphaned index attempts as failed
+    with get_session_with_default_tenant() as db_session:
+        unfenced_attempt_ids = get_unfenced_index_attempt_ids(db_session, r)
+        for attempt_id in unfenced_attempt_ids:
+            attempt = get_index_attempt(db_session, attempt_id)
+            if not attempt:
+                continue
+
+            failure_reason = (
+                f"Orphaned index attempt found on startup: "
+                f"index_attempt={attempt.id} "
+                f"cc_pair={attempt.connector_credential_pair_id} "
+                f"search_settings={attempt.search_settings_id}"
+            )
+            logger.warning(failure_reason)
+            mark_attempt_failed(attempt.id, db_session, failure_reason)
+
 
 @worker_ready.connect
 def on_worker_ready(sender: Any, **kwargs: Any) -> None:

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

@@ -1,12 +1,12 @@
 from datetime import datetime
 from datetime import timezone
 
-import redis
 from celery import Celery
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from redis import Redis
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.background.celery.apps.app_base import task_logger
@@ -87,7 +87,7 @@ def try_generate_document_cc_pair_cleanup_tasks(
     cc_pair_id: int,
     db_session: Session,
     r: Redis,
-    lock_beat: redis.lock.Lock,
+    lock_beat: RedisLock,
     tenant_id: str | None,
 ) -> int | None:
     """Returns an int if syncing is needed. The int represents the number of sync tasks generated.

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

@@ -3,13 +3,14 @@ from datetime import timezone
 from http import HTTPStatus
 from time import sleep
 
-import redis
 import sentry_sdk
 from celery import Celery
 from celery import shared_task
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
 from redis import Redis
+from redis.exceptions import LockError
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.background.celery.apps.app_base import task_logger
@@ -44,7 +45,7 @@ from danswer.db.swap_index import check_index_swap
 from danswer.natural_language_processing.search_nlp_models import EmbeddingModel
 from danswer.natural_language_processing.search_nlp_models import warm_up_bi_encoder
 from danswer.redis.redis_connector import RedisConnector
-from danswer.redis.redis_connector_index import RedisConnectorIndexingFenceData
+from danswer.redis.redis_connector_index import RedisConnectorIndexPayload
 from danswer.redis.redis_pool import get_redis_client
 from danswer.utils.logger import setup_logger
 from danswer.utils.variable_functionality import global_version
@@ -61,14 +62,18 @@ class RunIndexingCallback(RunIndexingCallbackInterface):
         self,
         stop_key: str,
         generator_progress_key: str,
-        redis_lock: redis.lock.Lock,
+        redis_lock: RedisLock,
         redis_client: Redis,
     ):
         super().__init__()
-        self.redis_lock: redis.lock.Lock = redis_lock
+        self.redis_lock: RedisLock = redis_lock
         self.stop_key: str = stop_key
         self.generator_progress_key: str = generator_progress_key
         self.redis_client = redis_client
+        self.started: datetime = datetime.now(timezone.utc)
+        self.redis_lock.reacquire()
+
+        self.last_lock_reacquire: datetime = datetime.now(timezone.utc)
 
     def should_stop(self) -> bool:
         if self.redis_client.exists(self.stop_key):
@@ -76,7 +81,19 @@ class RunIndexingCallback(RunIndexingCallbackInterface):
         return False
 
     def progress(self, amount: int) -> None:
-        self.redis_lock.reacquire()
+        try:
+            self.redis_lock.reacquire()
+            self.last_lock_reacquire = datetime.now(timezone.utc)
+        except LockError:
+            logger.exception(
+                f"RunIndexingCallback - lock.reacquire exceptioned. "
+                f"lock_timeout={self.redis_lock.timeout} "
+                f"start={self.started} "
+                f"last_reacquired={self.last_lock_reacquire} "
+                f"now={datetime.now(timezone.utc)}"
+            )
+            raise
+
         self.redis_client.incrby(self.generator_progress_key, amount)
 
 
@@ -325,7 +342,7 @@ def try_creating_indexing_task(
         redis_connector_index.generator_clear()
 
         # set a basic fence to start
-        payload = RedisConnectorIndexingFenceData(
+        payload = RedisConnectorIndexPayload(
             index_attempt_id=None,
             started=None,
             submitted=datetime.now(timezone.utc),
@@ -368,7 +385,7 @@ def try_creating_indexing_task(
         redis_connector_index.set_fence(payload)
 
     except Exception:
-        redis_connector_index.set_fence(payload)
+        redis_connector_index.set_fence(None)
         task_logger.exception(
             f"Unexpected exception: "
             f"tenant={tenant_id} "

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

@@ -13,6 +13,7 @@ from celery.exceptions import SoftTimeLimitExceeded
 from celery.result import AsyncResult
 from celery.states import READY_STATES
 from redis import Redis
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 from tenacity import RetryError
 
@@ -162,7 +163,7 @@ def try_generate_stale_document_sync_tasks(
     celery_app: Celery,
     db_session: Session,
     r: Redis,
-    lock_beat: redis.lock.Lock,
+    lock_beat: RedisLock,
     tenant_id: str | None,
 ) -> int | None:
     # the fence is up, do nothing
@@ -180,7 +181,12 @@ def try_generate_stale_document_sync_tasks(
         f"Stale documents found (at least {stale_doc_count}). Generating sync tasks by cc pair."
     )
 
-    task_logger.info("RedisConnector.generate_tasks starting by cc_pair.")
+    task_logger.info(
+        "RedisConnector.generate_tasks starting by cc_pair. "
+        "Documents spanning multiple cc_pairs will only be synced once."
+    )
+
+    docs_to_skip: set[str] = set()
 
     # rkuo: we could technically sync all stale docs in one big pass.
     # but I feel it's more understandable to group the docs by cc_pair
@@ -188,22 +194,21 @@ def try_generate_stale_document_sync_tasks(
     cc_pairs = get_connector_credential_pairs(db_session)
     for cc_pair in cc_pairs:
         rc = RedisConnectorCredentialPair(tenant_id, cc_pair.id)
-        tasks_generated = rc.generate_tasks(
-            celery_app, db_session, r, lock_beat, tenant_id
-        )
+        rc.set_skip_docs(docs_to_skip)
+        result = rc.generate_tasks(celery_app, db_session, r, lock_beat, tenant_id)
 
-        if tasks_generated is None:
+        if result is None:
             continue
 
-        if tasks_generated == 0:
+        if result[1] == 0:
             continue
 
         task_logger.info(
             f"RedisConnector.generate_tasks finished for single cc_pair. "
-            f"cc_pair_id={cc_pair.id} tasks_generated={tasks_generated}"
+            f"cc_pair={cc_pair.id} tasks_generated={result[0]} tasks_possible={result[1]}"
         )
 
-        total_tasks_generated += tasks_generated
+        total_tasks_generated += result[0]
 
     task_logger.info(
         f"RedisConnector.generate_tasks finished for all cc_pairs. total_tasks_generated={total_tasks_generated}"
@@ -218,7 +223,7 @@ def try_generate_document_set_sync_tasks(
     document_set_id: int,
     db_session: Session,
     r: Redis,
-    lock_beat: redis.lock.Lock,
+    lock_beat: RedisLock,
     tenant_id: str | None,
 ) -> int | None:
     lock_beat.reacquire()
@@ -246,12 +251,11 @@ def try_generate_document_set_sync_tasks(
     )
 
     # Add all documents that need to be updated into the queue
-    tasks_generated = rds.generate_tasks(
-        celery_app, db_session, r, lock_beat, tenant_id
-    )
-    if tasks_generated is None:
+    result = rds.generate_tasks(celery_app, db_session, r, lock_beat, tenant_id)
+    if result is None:
         return None
 
+    tasks_generated = result[0]
     # Currently we are allowing the sync to proceed with 0 tasks.
     # It's possible for sets/groups to be generated initially with no entries
     # and they still need to be marked as up to date.
@@ -260,7 +264,7 @@ def try_generate_document_set_sync_tasks(
 
     task_logger.info(
         f"RedisDocumentSet.generate_tasks finished. "
-        f"document_set_id={document_set.id} tasks_generated={tasks_generated}"
+        f"document_set={document_set.id} tasks_generated={tasks_generated}"
     )
 
     # set this only after all tasks have been added
@@ -273,7 +277,7 @@ def try_generate_user_group_sync_tasks(
     usergroup_id: int,
     db_session: Session,
     r: Redis,
-    lock_beat: redis.lock.Lock,
+    lock_beat: RedisLock,
     tenant_id: str | None,
 ) -> int | None:
     lock_beat.reacquire()
@@ -302,12 +306,11 @@ def try_generate_user_group_sync_tasks(
     task_logger.info(
         f"RedisUserGroup.generate_tasks starting. usergroup_id={usergroup.id}"
     )
-    tasks_generated = rug.generate_tasks(
-        celery_app, db_session, r, lock_beat, tenant_id
-    )
-    if tasks_generated is None:
+    result = rug.generate_tasks(celery_app, db_session, r, lock_beat, tenant_id)
+    if result is None:
         return None
 
+    tasks_generated = result[0]
     # Currently we are allowing the sync to proceed with 0 tasks.
     # It's possible for sets/groups to be generated initially with no entries
     # and they still need to be marked as up to date.
@@ -316,7 +319,7 @@ def try_generate_user_group_sync_tasks(
 
     task_logger.info(
         f"RedisUserGroup.generate_tasks finished. "
-        f"usergroup_id={usergroup.id} tasks_generated={tasks_generated}"
+        f"usergroup={usergroup.id} tasks_generated={tasks_generated}"
     )
 
     # set this only after all tasks have been added
@@ -580,8 +583,8 @@ def monitor_ccpair_indexing_taskset(
     progress = redis_connector_index.get_progress()
     if progress is not None:
         task_logger.info(
-            f"Connector indexing progress: cc_pair_id={cc_pair_id} "
-            f"search_settings_id={search_settings_id} "
+            f"Connector indexing progress: cc_pair={cc_pair_id} "
+            f"search_settings={search_settings_id} "
             f"progress={progress} "
             f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
         )
@@ -602,8 +605,8 @@ def monitor_ccpair_indexing_taskset(
             # if it isn't, then the worker crashed
             task_logger.info(
                 f"Connector indexing aborted: "
-                f"cc_pair_id={cc_pair_id} "
-                f"search_settings_id={search_settings_id} "
+                f"cc_pair={cc_pair_id} "
+                f"search_settings={search_settings_id} "
                 f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
             )
 
@@ -621,8 +624,8 @@ def monitor_ccpair_indexing_taskset(
     status_enum = HTTPStatus(status_int)
 
     task_logger.info(
-        f"Connector indexing finished: cc_pair_id={cc_pair_id} "
-        f"search_settings_id={search_settings_id} "
+        f"Connector indexing finished: cc_pair={cc_pair_id} "
+        f"search_settings={search_settings_id} "
         f"status={status_enum.name} "
         f"elapsed_submitted={elapsed_submitted.total_seconds():.2f}"
     )
@@ -630,6 +633,37 @@ def monitor_ccpair_indexing_taskset(
     redis_connector_index.reset()
 
 
+def get_unfenced_index_attempt_ids(db_session: Session, r: redis.Redis) -> list[int]:
+    """Gets a list of unfenced index attempts. Should not be possible, so we'd typically
+    want to clean them up.
+
+    Unfenced = attempt not in terminal state and fence does not exist.
+    """
+    unfenced_attempts: list[int] = []
+
+    # do some cleanup before clearing fences
+    # check the db for any outstanding index attempts
+    attempts: list[IndexAttempt] = []
+    attempts.extend(
+        get_all_index_attempts_by_status(IndexingStatus.NOT_STARTED, db_session)
+    )
+    attempts.extend(
+        get_all_index_attempts_by_status(IndexingStatus.IN_PROGRESS, db_session)
+    )
+
+    for attempt in attempts:
+        # if attempts exist in the db but we don't detect them in redis, mark them as failed
+        fence_key = RedisConnectorIndex.fence_key_with_ids(
+            attempt.connector_credential_pair_id, attempt.search_settings_id
+        )
+        if r.exists(fence_key):
+            continue
+
+        unfenced_attempts.append(attempt.id)
+
+    return unfenced_attempts
+
+
 @shared_task(name="monitor_vespa_sync", soft_time_limit=300, bind=True)
 def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     """This is a celery beat task that monitors and finalizes metadata sync tasksets.
@@ -643,7 +677,7 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
     """
     r = get_redis_client(tenant_id=tenant_id)
 
-    lock_beat: redis.lock.Lock = r.lock(
+    lock_beat: RedisLock = r.lock(
         DanswerRedisLocks.MONITOR_VESPA_SYNC_BEAT_LOCK,
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
@@ -677,31 +711,24 @@ def monitor_vespa_sync(self: Task, tenant_id: str | None) -> bool:
             f"pruning={n_pruning}"
         )
 
-        # do some cleanup before clearing fences
-        # check the db for any outstanding index attempts
+        # Fail any index attempts in the DB that don't have fences
         with get_session_with_tenant(tenant_id) as db_session:
-            attempts: list[IndexAttempt] = []
-            attempts.extend(
-                get_all_index_attempts_by_status(IndexingStatus.NOT_STARTED, db_session)
-            )
-            attempts.extend(
-                get_all_index_attempts_by_status(IndexingStatus.IN_PROGRESS, db_session)
-            )
+            unfenced_attempt_ids = get_unfenced_index_attempt_ids(db_session, r)
+            for attempt_id in unfenced_attempt_ids:
+                attempt = get_index_attempt(db_session, attempt_id)
+                if not attempt:
+                    continue
 
-            for a in attempts:
-                # if attempts exist in the db but we don't detect them in redis, mark them as failed
-                fence_key = RedisConnectorIndex.fence_key_with_ids(
-                    a.connector_credential_pair_id, a.search_settings_id
+                failure_reason = (
+                    f"Unfenced index attempt found in DB: "
+                    f"index_attempt={attempt.id} "
+                    f"cc_pair={attempt.connector_credential_pair_id} "
+                    f"search_settings={attempt.search_settings_id}"
+                )
+                task_logger.warning(failure_reason)
+                mark_attempt_failed(
+                    attempt.id, db_session, failure_reason=failure_reason
                 )
-                if not r.exists(fence_key):
-                    failure_reason = (
-                        f"Unknown index attempt. Might be left over from a process restart: "
-                        f"index_attempt={a.id} "
-                        f"cc_pair={a.connector_credential_pair_id} "
-                        f"search_settings={a.search_settings_id}"
-                    )
-                    task_logger.warning(failure_reason)
-                    mark_attempt_failed(a.id, db_session, failure_reason=failure_reason)
 
         lock_beat.reacquire()
         if r.exists(RedisConnectorCredentialPair.get_fence_key()):

backend/danswer/background/indexing/run_indexing.py

@@ -433,11 +433,13 @@ def run_indexing_entrypoint(
         with get_session_with_tenant(tenant_id) as db_session:
             attempt = transition_attempt_to_in_progress(index_attempt_id, db_session)
 
+            tenant_str = ""
+            if tenant_id is not None:
+                tenant_str = f" for tenant {tenant_id}"
+
             logger.info(
-                f"Indexing starting for tenant {tenant_id}: "
-                if tenant_id is not None
-                else ""
-                + f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"Indexing starting{tenant_str}: "
+                f"connector='{attempt.connector_credential_pair.connector.name}' "
                 f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )
@@ -445,10 +447,8 @@ def run_indexing_entrypoint(
             _run_indexing(db_session, attempt, tenant_id, callback)
 
             logger.info(
-                f"Indexing finished for tenant {tenant_id}: "
-                if tenant_id is not None
-                else ""
-                + f"connector='{attempt.connector_credential_pair.connector.name}' "
+                f"Indexing finished{tenant_str}: "
+                f"connector='{attempt.connector_credential_pair.connector.name}' "
                 f"config='{attempt.connector_credential_pair.connector.connector_specific_config}' "
                 f"credentials='{attempt.connector_credential_pair.connector_id}'"
             )

backend/danswer/chat/process_message.py

@@ -19,16 +19,10 @@ from danswer.chat.models import MessageSpecificCitations
 from danswer.chat.models import QADocsResponse
 from danswer.chat.models import StreamingError
 from danswer.chat.models import StreamStopInfo
-from danswer.configs.app_configs import AZURE_DALLE_API_BASE
-from danswer.configs.app_configs import AZURE_DALLE_API_KEY
-from danswer.configs.app_configs import AZURE_DALLE_API_VERSION
-from danswer.configs.app_configs import AZURE_DALLE_DEPLOYMENT_NAME
-from danswer.configs.chat_configs import BING_API_KEY
 from danswer.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from danswer.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
 from danswer.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT
 from danswer.configs.constants import MessageType
-from danswer.configs.model_configs import GEN_AI_TEMPERATURE
 from danswer.db.chat import attach_files_to_chat_message
 from danswer.db.chat import create_db_search_doc
 from danswer.db.chat import create_new_chat_message
@@ -41,7 +35,6 @@ from danswer.db.chat import reserve_message_id
 from danswer.db.chat import translate_db_message_to_chat_message_detail
 from danswer.db.chat import translate_db_search_doc_to_server_search_doc
 from danswer.db.engine import get_session_context_manager
-from danswer.db.llm import fetch_existing_llm_providers
 from danswer.db.models import SearchDoc as DbSearchDoc
 from danswer.db.models import ToolCall
 from danswer.db.models import User
@@ -61,14 +54,13 @@ from danswer.llm.answering.models import PromptConfig
 from danswer.llm.exceptions import GenAIDisabledException
 from danswer.llm.factory import get_llms_for_persona
 from danswer.llm.factory import get_main_llm_from_tuple
-from danswer.llm.interfaces import LLMConfig
 from danswer.llm.utils import litellm_exception_to_error_msg
 from danswer.natural_language_processing.utils import get_tokenizer
-from danswer.search.enums import LLMEvaluationType
 from danswer.search.enums import OptionalSearchSetting
 from danswer.search.enums import QueryFlow
 from danswer.search.enums import SearchType
 from danswer.search.models import InferenceSection
+from danswer.search.models import RetrievalDetails
 from danswer.search.retrieval.search_runner import inference_sections_from_ids
 from danswer.search.utils import chunks_or_sections_to_search_docs
 from danswer.search.utils import dedupe_documents
@@ -77,14 +69,14 @@ from danswer.search.utils import relevant_sections_to_indices
 from danswer.server.query_and_chat.models import ChatMessageDetail
 from danswer.server.query_and_chat.models import CreateChatMessageRequest
 from danswer.server.utils import get_json_line
-from danswer.tools.built_in_tools import get_built_in_tool_by_id
 from danswer.tools.force import ForceUseTool
-from danswer.tools.models import DynamicSchemaInfo
 from danswer.tools.models import ToolResponse
 from danswer.tools.tool import Tool
-from danswer.tools.tool_implementations.custom.custom_tool import (
-    build_custom_tools_from_openapi_schema_and_headers,
-)
+from danswer.tools.tool_constructor import construct_tools
+from danswer.tools.tool_constructor import CustomToolConfig
+from danswer.tools.tool_constructor import ImageGenerationToolConfig
+from danswer.tools.tool_constructor import InternetSearchToolConfig
+from danswer.tools.tool_constructor import SearchToolConfig
 from danswer.tools.tool_implementations.custom.custom_tool import (
     CUSTOM_TOOL_RESPONSE_ID,
 )
@@ -95,9 +87,6 @@ from danswer.tools.tool_implementations.images.image_generation_tool import (
 from danswer.tools.tool_implementations.images.image_generation_tool import (
     ImageGenerationResponse,
 )
-from danswer.tools.tool_implementations.images.image_generation_tool import (
-    ImageGenerationTool,
-)
 from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
     INTERNET_SEARCH_RESPONSE_ID,
 )
@@ -122,9 +111,6 @@ from danswer.tools.tool_implementations.search.search_tool import (
     SECTION_RELEVANCE_LIST_ID,
 )
 from danswer.tools.tool_runner import ToolCallFinalResult
-from danswer.tools.utils import compute_all_tool_tokens
-from danswer.tools.utils import explicit_tool_calling_supported
-from danswer.utils.headers import header_dict_to_header_list
 from danswer.utils.logger import setup_logger
 from danswer.utils.timing import log_generator_function_time
 
@@ -295,7 +281,6 @@ def stream_chat_message_objects(
     max_document_percentage: float = CHAT_TARGET_CHUNK_PERCENTAGE,
     # if specified, uses the last user message and does not create a new user message based
     # on the `new_msg_req.message`. Currently, requires a state where the last message is a
-    use_existing_user_message: bool = False,
     litellm_additional_headers: dict[str, str] | None = None,
     custom_tool_additional_headers: dict[str, str] | None = None,
     is_connected: Callable[[], bool] | None = None,
@@ -307,6 +292,9 @@ def stream_chat_message_objects(
     3. [always] A set of streamed LLM tokens or an error anywhere along the line if something fails
     4. [always] Details on the final AI response message that is created
     """
+    use_existing_user_message = new_msg_req.use_existing_user_message
+    existing_assistant_message_id = new_msg_req.existing_assistant_message_id
+
     # Currently surrounding context is not supported for chat
     # Chat is already token heavy and harder for the model to process plus it would roll history over much faster
     new_msg_req.chunks_above = 0
@@ -428,12 +416,20 @@ def stream_chat_message_objects(
             final_msg, history_msgs = create_chat_chain(
                 chat_session_id=chat_session_id, db_session=db_session
             )
-            if final_msg.message_type != MessageType.USER:
-                raise RuntimeError(
-                    "The last message was not a user message. Cannot call "
-                    "`stream_chat_message_objects` with `is_regenerate=True` "
-                    "when the last message is not a user message."
-                )
+            if existing_assistant_message_id is None:
+                if final_msg.message_type != MessageType.USER:
+                    raise RuntimeError(
+                        "The last message was not a user message. Cannot call "
+                        "`stream_chat_message_objects` with `is_regenerate=True` "
+                        "when the last message is not a user message."
+                    )
+            else:
+                if final_msg.id != existing_assistant_message_id:
+                    raise RuntimeError(
+                        "The last message was not the existing assistant message. "
+                        f"Final message id: {final_msg.id}, "
+                        f"existing assistant message id: {existing_assistant_message_id}"
+                    )
 
         # Disable Query Rephrasing for the first message
         # This leads to a better first response since the LLM rephrasing the question
@@ -504,13 +500,19 @@ def stream_chat_message_objects(
                 ),
                 max_window_percentage=max_document_percentage,
             )
-        reserved_message_id = reserve_message_id(
-            db_session=db_session,
-            chat_session_id=chat_session_id,
-            parent_message=user_message.id
-            if user_message is not None
-            else parent_message.id,
-            message_type=MessageType.ASSISTANT,
+
+        # we don't need to reserve a message id if we're using an existing assistant message
+        reserved_message_id = (
+            final_msg.id
+            if existing_assistant_message_id is not None
+            else reserve_message_id(
+                db_session=db_session,
+                chat_session_id=chat_session_id,
+                parent_message=user_message.id
+                if user_message is not None
+                else parent_message.id,
+                message_type=MessageType.ASSISTANT,
+            )
         )
         yield MessageResponseIDInfo(
             user_message_id=user_message.id if user_message else None,
@@ -525,7 +527,13 @@ def stream_chat_message_objects(
         partial_response = partial(
             create_new_chat_message,
             chat_session_id=chat_session_id,
-            parent_message=final_msg,
+            # if we're using an existing assistant message, then this will just be an
+            # update operation, in which case the parent should be the parent of
+            # the latest. If we're creating a new assistant message, then the parent
+            # should be the latest message (latest user message)
+            parent_message=(
+                final_msg if existing_assistant_message_id is None else parent_message
+            ),
             prompt_id=prompt_id,
             overridden_model=overridden_model,
             # message=,
@@ -537,6 +545,7 @@ def stream_chat_message_objects(
             # reference_docs=,
             db_session=db_session,
             commit=False,
+            reserved_message_id=reserved_message_id,
         )
 
         if not final_msg.prompt:
@@ -560,142 +569,39 @@ def stream_chat_message_objects(
             structured_response_format=new_msg_req.structured_response_format,
         )
 
-        # find out what tools to use
-        search_tool: SearchTool | None = None
-        tool_dict: dict[int, list[Tool]] = {}  # tool_id to tool
-        for db_tool_model in persona.tools:
-            # handle in-code tools specially
-            if db_tool_model.in_code_tool_id:
-                tool_cls = get_built_in_tool_by_id(db_tool_model.id, db_session)
-                if tool_cls.__name__ == SearchTool.__name__ and not latest_query_files:
-                    search_tool = SearchTool(
-                        db_session=db_session,
-                        user=user,
-                        persona=persona,
-                        retrieval_options=retrieval_options,
-                        prompt_config=prompt_config,
-                        llm=llm,
-                        fast_llm=fast_llm,
-                        pruning_config=document_pruning_config,
-                        answer_style_config=answer_style_config,
-                        selected_sections=selected_sections,
-                        chunks_above=new_msg_req.chunks_above,
-                        chunks_below=new_msg_req.chunks_below,
-                        full_doc=new_msg_req.full_doc,
-                        evaluation_type=(
-                            LLMEvaluationType.BASIC
-                            if persona.llm_relevance_filter
-                            else LLMEvaluationType.SKIP
-                        ),
-                    )
-                    tool_dict[db_tool_model.id] = [search_tool]
-                elif tool_cls.__name__ == ImageGenerationTool.__name__:
-                    img_generation_llm_config: LLMConfig | None = None
-                    if (
-                        llm
-                        and llm.config.api_key
-                        and llm.config.model_provider == "openai"
-                    ):
-                        img_generation_llm_config = LLMConfig(
-                            model_provider=llm.config.model_provider,
-                            model_name="dall-e-3",
-                            temperature=GEN_AI_TEMPERATURE,
-                            api_key=llm.config.api_key,
-                            api_base=llm.config.api_base,
-                            api_version=llm.config.api_version,
-                        )
-                    elif (
-                        llm.config.model_provider == "azure"
-                        and AZURE_DALLE_API_KEY is not None
-                    ):
-                        img_generation_llm_config = LLMConfig(
-                            model_provider="azure",
-                            model_name=f"azure/{AZURE_DALLE_DEPLOYMENT_NAME}",
-                            temperature=GEN_AI_TEMPERATURE,
-                            api_key=AZURE_DALLE_API_KEY,
-                            api_base=AZURE_DALLE_API_BASE,
-                            api_version=AZURE_DALLE_API_VERSION,
-                        )
-                    else:
-                        llm_providers = fetch_existing_llm_providers(db_session)
-                        openai_provider = next(
-                            iter(
-                                [
-                                    llm_provider
-                                    for llm_provider in llm_providers
-                                    if llm_provider.provider == "openai"
-                                ]
-                            ),
-                            None,
-                        )
-                        if not openai_provider or not openai_provider.api_key:
-                            raise ValueError(
-                                "Image generation tool requires an OpenAI API key"
-                            )
-                        img_generation_llm_config = LLMConfig(
-                            model_provider=openai_provider.provider,
-                            model_name="dall-e-3",
-                            temperature=GEN_AI_TEMPERATURE,
-                            api_key=openai_provider.api_key,
-                            api_base=openai_provider.api_base,
-                            api_version=openai_provider.api_version,
-                        )
-                    tool_dict[db_tool_model.id] = [
-                        ImageGenerationTool(
-                            api_key=cast(str, img_generation_llm_config.api_key),
-                            api_base=img_generation_llm_config.api_base,
-                            api_version=img_generation_llm_config.api_version,
-                            additional_headers=litellm_additional_headers,
-                            model=img_generation_llm_config.model_name,
-                        )
-                    ]
-                elif tool_cls.__name__ == InternetSearchTool.__name__:
-                    bing_api_key = BING_API_KEY
-                    if not bing_api_key:
-                        raise ValueError(
-                            "Internet search tool requires a Bing API key, please contact your Danswer admin to get it added!"
-                        )
-                    tool_dict[db_tool_model.id] = [
-                        InternetSearchTool(
-                            api_key=bing_api_key,
-                            answer_style_config=answer_style_config,
-                            prompt_config=prompt_config,
-                        )
-                    ]
-
-                continue
-
-            # handle all custom tools
-            if db_tool_model.openapi_schema:
-                tool_dict[db_tool_model.id] = cast(
-                    list[Tool],
-                    build_custom_tools_from_openapi_schema_and_headers(
-                        db_tool_model.openapi_schema,
-                        dynamic_schema_info=DynamicSchemaInfo(
-                            chat_session_id=chat_session_id,
-                            message_id=user_message.id if user_message else None,
-                        ),
-                        custom_headers=(db_tool_model.custom_headers or [])
-                        + (
-                            header_dict_to_header_list(
-                                custom_tool_additional_headers or {}
-                            )
-                        ),
-                    ),
-                )
-
+        tool_dict = construct_tools(
+            persona=persona,
+            prompt_config=prompt_config,
+            db_session=db_session,
+            user=user,
+            llm=llm,
+            fast_llm=fast_llm,
+            search_tool_config=SearchToolConfig(
+                answer_style_config=answer_style_config,
+                document_pruning_config=document_pruning_config,
+                retrieval_options=retrieval_options or RetrievalDetails(),
+                selected_sections=selected_sections,
+                chunks_above=new_msg_req.chunks_above,
+                chunks_below=new_msg_req.chunks_below,
+                full_doc=new_msg_req.full_doc,
+                latest_query_files=latest_query_files,
+            ),
+            internet_search_tool_config=InternetSearchToolConfig(
+                answer_style_config=answer_style_config,
+            ),
+            image_generation_tool_config=ImageGenerationToolConfig(
+                additional_headers=litellm_additional_headers,
+            ),
+            custom_tool_config=CustomToolConfig(
+                chat_session_id=chat_session_id,
+                message_id=user_message.id if user_message else None,
+                additional_headers=custom_tool_additional_headers,
+            ),
+        )
         tools: list[Tool] = []
         for tool_list in tool_dict.values():
             tools.extend(tool_list)
 
-        # factor in tool definition size when pruning
-        document_pruning_config.tool_num_tokens = compute_all_tool_tokens(
-            tools, llm_tokenizer
-        )
-        document_pruning_config.using_tool_message = explicit_tool_calling_supported(
-            llm_provider, llm_model_name
-        )
-
         # LLM prompt building, response capturing, etc.
         answer = Answer(
             is_connected=is_connected,
@@ -871,7 +777,6 @@ def stream_chat_message_objects(
                 tool_name_to_tool_id[tool.name] = tool_id
 
         gen_ai_response_message = partial_response(
-            reserved_message_id=reserved_message_id,
             message=answer.llm_answer,
             rephrased_query=(
                 qa_docs_response.rephrased_query if qa_docs_response else None
@@ -879,9 +784,11 @@ def stream_chat_message_objects(
             reference_docs=reference_db_search_docs,
             files=ai_message_files,
             token_count=len(llm_tokenizer_encode_func(answer.llm_answer)),
-            citations=message_specific_citations.citation_map
-            if message_specific_citations
-            else None,
+            citations=(
+                message_specific_citations.citation_map
+                if message_specific_citations
+                else None
+            ),
             error=None,
             tool_call=(
                 ToolCall(
@@ -915,7 +822,6 @@ def stream_chat_message_objects(
 def stream_chat_message(
     new_msg_req: CreateChatMessageRequest,
     user: User | None,
-    use_existing_user_message: bool = False,
     litellm_additional_headers: dict[str, str] | None = None,
     custom_tool_additional_headers: dict[str, str] | None = None,
     is_connected: Callable[[], bool] | None = None,
@@ -925,7 +831,6 @@ def stream_chat_message(
             new_msg_req=new_msg_req,
             user=user,
             db_session=db_session,
-            use_existing_user_message=use_existing_user_message,
             litellm_additional_headers=litellm_additional_headers,
             custom_tool_additional_headers=custom_tool_additional_headers,
             is_connected=is_connected,

backend/danswer/configs/constants.py

@@ -74,7 +74,7 @@ CELERY_PRIMARY_WORKER_LOCK_TIMEOUT = 120
 
 # needs to be long enough to cover the maximum time it takes to download an object
 # if we can get callbacks as object bytes download, we could lower this a lot.
-CELERY_INDEXING_LOCK_TIMEOUT = 60 * 60  # 60 min
+CELERY_INDEXING_LOCK_TIMEOUT = 3 * 60 * 60  # 60 min
 
 # needs to be long enough to cover the maximum time it takes to download an object
 # if we can get callbacks as object bytes download, we could lower this a lot.

backend/danswer/db/document.py

@@ -169,6 +169,7 @@ def get_document_connector_counts(
 def get_document_counts_for_cc_pairs(
     db_session: Session, cc_pair_identifiers: list[ConnectorCredentialPairIdentifier]
 ) -> Sequence[tuple[int, int, int]]:
+    """Returns a sequence of tuples of (connector_id, credential_id, document count)"""
     stmt = (
         select(
             DocumentByConnectorCredentialPair.connector_id,
@@ -323,23 +324,23 @@ def upsert_documents(
 
 
 def upsert_document_by_connector_credential_pair(
-    db_session: Session, document_metadata_batch: list[DocumentMetadata]
+    db_session: Session, connector_id: int, credential_id: int, document_ids: list[str]
 ) -> None:
     """NOTE: this function is Postgres specific. Not all DBs support the ON CONFLICT clause."""
-    if not document_metadata_batch:
-        logger.info("`document_metadata_batch` is empty. Skipping.")
+    if not document_ids:
+        logger.info("`document_ids` is empty. Skipping.")
         return
 
     insert_stmt = insert(DocumentByConnectorCredentialPair).values(
         [
             model_to_dict(
                 DocumentByConnectorCredentialPair(
-                    id=document_metadata.document_id,
-                    connector_id=document_metadata.connector_id,
-                    credential_id=document_metadata.credential_id,
+                    id=doc_id,
+                    connector_id=connector_id,
+                    credential_id=credential_id,
                 )
             )
-            for document_metadata in document_metadata_batch
+            for doc_id in document_ids
         ]
     )
     # for now, there are no columns to update. If more metadata is added, then this
@@ -400,17 +401,6 @@ def mark_document_as_synced(document_id: str, db_session: Session) -> None:
     db_session.commit()
 
 
-def upsert_documents_complete(
-    db_session: Session,
-    document_metadata_batch: list[DocumentMetadata],
-) -> None:
-    upsert_documents(db_session, document_metadata_batch)
-    upsert_document_by_connector_credential_pair(db_session, document_metadata_batch)
-    logger.info(
-        f"Upserted {len(document_metadata_batch)} document store entries into DB"
-    )
-
-
 def delete_document_by_connector_credential_pair__no_commit(
     db_session: Session,
     document_id: str,
@@ -520,7 +510,7 @@ def prepare_to_modify_documents(
     db_session.commit()  # ensure that we're not in a transaction
 
     lock_acquired = False
-    for _ in range(_NUM_LOCK_ATTEMPTS):
+    for i in range(_NUM_LOCK_ATTEMPTS):
         try:
             with db_session.begin() as transaction:
                 lock_acquired = acquire_document_locks(
@@ -531,7 +521,7 @@ def prepare_to_modify_documents(
                     break
         except OperationalError as e:
             logger.warning(
-                f"Failed to acquire locks for documents, retrying. Error: {e}"
+                f"Failed to acquire locks for documents on attempt {i}, retrying. Error: {e}"
             )
 
         time.sleep(retry_delay)

backend/danswer/db/engine.py

@@ -312,7 +312,9 @@ async def get_async_session_with_tenant(
             await session.execute(text(f'SET search_path = "{tenant_id}"'))
             if POSTGRES_IDLE_SESSIONS_TIMEOUT:
                 await session.execute(
-                    f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                    text(
+                        f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                    )
                 )
         except Exception:
             logger.exception("Error setting search_path.")
@@ -373,7 +375,9 @@ def get_session_with_tenant(
                 cursor.execute(f'SET search_path = "{tenant_id}"')
                 if POSTGRES_IDLE_SESSIONS_TIMEOUT:
                     cursor.execute(
-                        f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                        text(
+                            f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                        )
                     )
             finally:
                 cursor.close()

backend/danswer/db/tools.py

@@ -24,6 +24,13 @@ def get_tool_by_id(tool_id: int, db_session: Session) -> Tool:
     return tool
 
 
+def get_tool_by_name(tool_name: str, db_session: Session) -> Tool:
+    tool = db_session.scalar(select(Tool).where(Tool.name == tool_name))
+    if not tool:
+        raise ValueError("Tool by specified name does not exist")
+    return tool
+
+
 def create_tool(
     name: str,
     description: str | None,
@@ -37,7 +44,7 @@ def create_tool(
         description=description,
         in_code_tool_id=None,
         openapi_schema=openapi_schema,
-        custom_headers=[header.dict() for header in custom_headers]
+        custom_headers=[header.model_dump() for header in custom_headers]
         if custom_headers
         else [],
         user_id=user_id,

backend/danswer/indexing/indexing_pipeline.py

@@ -20,7 +20,8 @@ from danswer.db.document import get_documents_by_ids
 from danswer.db.document import prepare_to_modify_documents
 from danswer.db.document import update_docs_last_modified__no_commit
 from danswer.db.document import update_docs_updated_at__no_commit
-from danswer.db.document import upsert_documents_complete
+from danswer.db.document import upsert_document_by_connector_credential_pair
+from danswer.db.document import upsert_documents
 from danswer.db.document_set import fetch_document_sets_for_documents
 from danswer.db.index_attempt import create_index_attempt_error
 from danswer.db.models import Document as DBDocument
@@ -56,13 +57,13 @@ class IndexingPipelineProtocol(Protocol):
         ...
 
 
-def upsert_documents_in_db(
+def _upsert_documents_in_db(
     documents: list[Document],
     index_attempt_metadata: IndexAttemptMetadata,
     db_session: Session,
 ) -> None:
     # Metadata here refers to basic document info, not metadata about the actual content
-    doc_m_batch: list[DocumentMetadata] = []
+    document_metadata_list: list[DocumentMetadata] = []
     for doc in documents:
         first_link = next(
             (section.link for section in doc.sections if section.link), ""
@@ -77,12 +78,9 @@ def upsert_documents_in_db(
             secondary_owners=get_experts_stores_representations(doc.secondary_owners),
             from_ingestion_api=doc.from_ingestion_api,
         )
-        doc_m_batch.append(db_doc_metadata)
+        document_metadata_list.append(db_doc_metadata)
 
-    upsert_documents_complete(
-        db_session=db_session,
-        document_metadata_batch=doc_m_batch,
-    )
+    upsert_documents(db_session, document_metadata_list)
 
     # Insert document content metadata
     for doc in documents:
@@ -95,21 +93,25 @@ def upsert_documents_in_db(
                     document_id=doc.id,
                     db_session=db_session,
                 )
-            else:
-                create_or_add_document_tag(
-                    tag_key=k,
-                    tag_value=v,
-                    source=doc.source,
-                    document_id=doc.id,
-                    db_session=db_session,
-                )
+                continue
+
+            create_or_add_document_tag(
+                tag_key=k,
+                tag_value=v,
+                source=doc.source,
+                document_id=doc.id,
+                db_session=db_session,
+            )
 
 
 def get_doc_ids_to_update(
     documents: list[Document], db_docs: list[DBDocument]
 ) -> list[Document]:
     """Figures out which documents actually need to be updated. If a document is already present
-    and the `updated_at` hasn't changed, we shouldn't need to do anything with it."""
+    and the `updated_at` hasn't changed, we shouldn't need to do anything with it.
+
+    NB: Still need to associate the document in the DB if multiple connectors are
+    indexing the same doc."""
     id_update_time_map = {
         doc.id: doc.doc_updated_at for doc in db_docs if doc.doc_updated_at
     }
@@ -195,9 +197,9 @@ def index_doc_batch_prepare(
     db_session: Session,
     ignore_time_skip: bool = False,
 ) -> DocumentBatchPrepareContext | None:
-    """This sets up the documents in the relational DB (source of truth) for permissions, metadata, etc.
+    """Sets up the documents in the relational DB (source of truth) for permissions, metadata, etc.
     This preceeds indexing it into the actual document index."""
-    documents = []
+    documents: list[Document] = []
     for document in document_batch:
         empty_contents = not any(section.text.strip() for section in document.sections)
         if (
@@ -212,43 +214,58 @@ def index_doc_batch_prepare(
             logger.warning(
                 f"Skipping document with ID {document.id} as it has neither title nor content."
             )
-        elif (
-            document.title is not None and not document.title.strip() and empty_contents
-        ):
+            continue
+
+        if document.title is not None and not document.title.strip() and empty_contents:
             # The title is explicitly empty ("" and not None) and the document is empty
             # so when building the chunk text representation, it will be empty and unuseable
             logger.warning(
                 f"Skipping document with ID {document.id} as the chunks will be empty."
             )
-        else:
-            documents.append(document)
+            continue
 
-    document_ids = [document.id for document in documents]
+        documents.append(document)
+
+    # Create a trimmed list of docs that don't have a newer updated at
+    # Shortcuts the time-consuming flow on connector index retries
+    document_ids: list[str] = [document.id for document in documents]
     db_docs: list[DBDocument] = get_documents_by_ids(
         db_session=db_session,
         document_ids=document_ids,
     )
 
-    # Skip indexing docs that don't have a newer updated at
-    # Shortcuts the time-consuming flow on connector index retries
     updatable_docs = (
         get_doc_ids_to_update(documents=documents, db_docs=db_docs)
         if not ignore_time_skip
         else documents
     )
 
-    # No docs to update either because the batch is empty or every doc was already indexed
+    # for all updatable docs, upsert into the DB
+    # Does not include doc_updated_at which is also used to indicate a successful update
+    if updatable_docs:
+        _upsert_documents_in_db(
+            documents=updatable_docs,
+            index_attempt_metadata=index_attempt_metadata,
+            db_session=db_session,
+        )
+
+    logger.info(
+        f"Upserted {len(updatable_docs)} changed docs out of "
+        f"{len(documents)} total docs into the DB"
+    )
+
+    # for all docs, upsert the document to cc pair relationship
+    upsert_document_by_connector_credential_pair(
+        db_session,
+        index_attempt_metadata.connector_id,
+        index_attempt_metadata.credential_id,
+        document_ids,
+    )
+
+    # No docs to process because the batch is empty or every doc was already indexed
     if not updatable_docs:
         return None
 
-    # Create records in the source of truth about these documents,
-    # does not include doc_updated_at which is also used to indicate a successful update
-    upsert_documents_in_db(
-        documents=documents,
-        index_attempt_metadata=index_attempt_metadata,
-        db_session=db_session,
-    )
-
     id_to_db_doc_map = {doc.id: doc for doc in db_docs}
     return DocumentBatchPrepareContext(
         updatable_docs=updatable_docs, id_to_db_doc_map=id_to_db_doc_map
@@ -269,7 +286,10 @@ def index_doc_batch(
 ) -> tuple[int, int]:
     """Takes different pieces of the indexing pipeline and applies it to a batch of documents
     Note that the documents should already be batched at this point so that it does not inflate the
-    memory requirements"""
+    memory requirements
+
+    Returns a tuple where the first element is the number of new docs and the
+    second element is the number of chunks."""
 
     no_access = DocumentAccess.build(
         user_emails=[],
@@ -312,9 +332,9 @@ def index_doc_batch(
 
         # we're concerned about race conditions where multiple simultaneous indexings might result
         # in one set of metadata overwriting another one in vespa.
-        # we still write data here for immediate and most likely correct sync, but
+        # we still write data here for the immediate and most likely correct sync, but
         # to resolve this, an update of the last modified field at the end of this loop
-        # always triggers a final metadata sync
+        # always triggers a final metadata sync via the celery queue
         access_aware_chunks = [
             DocMetadataAwareIndexChunk.from_index_chunk(
                 index_chunk=chunk,
@@ -351,7 +371,8 @@ def index_doc_batch(
         ids_to_new_updated_at = {}
         for doc in successful_docs:
             last_modified_ids.append(doc.id)
-            # doc_updated_at is the connector source's idea of when the doc was last modified
+            # doc_updated_at is the source's idea (on the other end of the connector)
+            # of when the doc was last modified
             if doc.doc_updated_at is None:
                 continue
             ids_to_new_updated_at[doc.id] = doc.doc_updated_at
@@ -366,10 +387,13 @@ def index_doc_batch(
 
         db_session.commit()
 
-    return len([r for r in insertion_records if r.already_existed is False]), len(
-        access_aware_chunks
+    result = (
+        len([r for r in insertion_records if r.already_existed is False]),
+        len(access_aware_chunks),
     )
 
+    return result
+
 
 def build_indexing_pipeline(
     *,

backend/danswer/main.py

@@ -25,6 +25,7 @@ from danswer.auth.schemas import UserCreate
 from danswer.auth.schemas import UserRead
 from danswer.auth.schemas import UserUpdate
 from danswer.auth.users import auth_backend
+from danswer.auth.users import BasicAuthenticationError
 from danswer.auth.users import fastapi_users
 from danswer.configs.app_configs import APP_API_PREFIX
 from danswer.configs.app_configs import APP_HOST
@@ -73,6 +74,9 @@ from danswer.server.manage.search_settings import router as search_settings_rout
 from danswer.server.manage.slack_bot import router as slack_bot_management_router
 from danswer.server.manage.users import router as user_router
 from danswer.server.middleware.latency_logging import add_latency_logging_middleware
+from danswer.server.openai_assistants_api.full_openai_assistants_api import (
+    get_full_openai_assistants_api_router,
+)
 from danswer.server.query_and_chat.chat_backend import router as chat_router
 from danswer.server.query_and_chat.query_backend import (
     admin_router as admin_query_router,
@@ -194,7 +198,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
 
 def log_http_error(_: Request, exc: Exception) -> JSONResponse:
     status_code = getattr(exc, "status_code", 500)
-    if status_code >= 400:
+
+    if isinstance(exc, BasicAuthenticationError):
+        # For BasicAuthenticationError, just log a brief message without stack trace (almost always spam)
+        logger.error(f"Authentication failed: {str(exc)}")
+
+    elif status_code >= 400:
         error_msg = f"{str(exc)}\n"
         error_msg += "".join(traceback.format_tb(exc.__traceback__))
         logger.error(error_msg)
@@ -220,7 +229,6 @@ def get_application() -> FastAPI:
     else:
         logger.debug("Sentry DSN not provided, skipping Sentry initialization")
 
-    # Add the custom exception handler
     application.add_exception_handler(status.HTTP_400_BAD_REQUEST, log_http_error)
     application.add_exception_handler(status.HTTP_401_UNAUTHORIZED, log_http_error)
     application.add_exception_handler(status.HTTP_403_FORBIDDEN, log_http_error)
@@ -265,6 +273,9 @@ def get_application() -> FastAPI:
         application, token_rate_limit_settings_router
     )
     include_router_with_global_prefix_prepended(application, indexing_router)
+    include_router_with_global_prefix_prepended(
+        application, get_full_openai_assistants_api_router()
+    )
 
     if AUTH_TYPE == AuthType.DISABLED:
         # Server logs this during auth setup verification step

backend/danswer/redis/redis_connector_credential_pair.py

@@ -1,9 +1,10 @@
 import time
+from typing import cast
 from uuid import uuid4
 
-import redis
 from celery import Celery
 from redis import Redis
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
@@ -13,6 +14,7 @@ from danswer.db.connector_credential_pair import get_connector_credential_pair_f
 from danswer.db.document import (
     construct_document_select_for_connector_credential_pair_by_needs_sync,
 )
+from danswer.db.models import Document
 from danswer.redis.redis_object_helper import RedisObjectHelper
 
 
@@ -30,6 +32,9 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
     def __init__(self, tenant_id: str | None, id: int) -> None:
         super().__init__(tenant_id, str(id))
 
+        # documents that should be skipped
+        self.skip_docs: set[str] = set()
+
     @classmethod
     def get_fence_key(cls) -> str:
         return RedisConnectorCredentialPair.FENCE_PREFIX
@@ -45,14 +50,19 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
         # example: connector_taskset
         return f"{self.TASKSET_PREFIX}"
 
+    def set_skip_docs(self, skip_docs: set[str]) -> None:
+        # documents that should be skipped. Note that this classes updates
+        # the list on the fly
+        self.skip_docs = skip_docs
+
     def generate_tasks(
         self,
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
-        lock: redis.lock.Lock,
+        lock: RedisLock,
         tenant_id: str | None,
-    ) -> int | None:
+    ) -> tuple[int, int] | None:
         last_lock_time = time.monotonic()
 
         async_results = []
@@ -63,7 +73,10 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
         stmt = construct_document_select_for_connector_credential_pair_by_needs_sync(
             cc_pair.connector_id, cc_pair.credential_id
         )
+
+        num_docs = 0
         for doc in db_session.scalars(stmt).yield_per(1):
+            doc = cast(Document, doc)
             current_time = time.monotonic()
             if current_time - last_lock_time >= (
                 CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4
@@ -71,6 +84,12 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
                 lock.reacquire()
                 last_lock_time = current_time
 
+            num_docs += 1
+
+            # check if we should skip the document (typically because it's already syncing)
+            if doc.id in self.skip_docs:
+                continue
+
             # celery's default task id format is "dd32ded3-00aa-4884-8b21-42f8332e7fac"
             # the key for the result is "celery-task-meta-dd32ded3-00aa-4884-8b21-42f8332e7fac"
             # we prefix the task id so it's easier to keep track of who created the task
@@ -93,5 +112,6 @@ class RedisConnectorCredentialPair(RedisObjectHelper):
             )
 
             async_results.append(result)
+            self.skip_docs.add(doc.id)
 
-        return len(async_results)
+        return len(async_results), num_docs

backend/danswer/redis/redis_connector_delete.py

@@ -6,6 +6,7 @@ from uuid import uuid4
 import redis
 from celery import Celery
 from pydantic import BaseModel
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
@@ -83,7 +84,7 @@ class RedisConnectorDelete:
         self,
         celery_app: Celery,
         db_session: Session,
-        lock: redis.lock.Lock,
+        lock: RedisLock,
     ) -> int | None:
         """Returns None if the cc_pair doesn't exist.
         Otherwise, returns an int with the number of generated tasks."""

backend/danswer/redis/redis_connector_index.py

@@ -6,7 +6,7 @@ import redis
 from pydantic import BaseModel
 
 
-class RedisConnectorIndexingFenceData(BaseModel):
+class RedisConnectorIndexPayload(BaseModel):
     index_attempt_id: int | None
     started: datetime | None
     submitted: datetime
@@ -71,22 +71,20 @@ class RedisConnectorIndex:
         return False
 
     @property
-    def payload(self) -> RedisConnectorIndexingFenceData | None:
+    def payload(self) -> RedisConnectorIndexPayload | None:
         # read related data and evaluate/print task progress
         fence_bytes = cast(bytes, self.redis.get(self.fence_key))
         if fence_bytes is None:
             return None
 
         fence_str = fence_bytes.decode("utf-8")
-        payload = RedisConnectorIndexingFenceData.model_validate_json(
-            cast(str, fence_str)
-        )
+        payload = RedisConnectorIndexPayload.model_validate_json(cast(str, fence_str))
 
         return payload
 
     def set_fence(
         self,
-        payload: RedisConnectorIndexingFenceData | None,
+        payload: RedisConnectorIndexPayload | None,
     ) -> None:
         if not payload:
             self.redis.delete(self.fence_key)

backend/danswer/redis/redis_connector_prune.py

@@ -4,6 +4,7 @@ from uuid import uuid4
 
 import redis
 from celery import Celery
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
@@ -105,7 +106,7 @@ class RedisConnectorPrune:
         documents_to_prune: set[str],
         celery_app: Celery,
         db_session: Session,
-        lock: redis.lock.Lock | None,
+        lock: RedisLock | None,
     ) -> int | None:
         last_lock_time = time.monotonic()
 

backend/danswer/redis/redis_document_set.py

@@ -5,6 +5,7 @@ from uuid import uuid4
 import redis
 from celery import Celery
 from redis import Redis
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
@@ -50,9 +51,9 @@ class RedisDocumentSet(RedisObjectHelper):
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
-        lock: redis.lock.Lock,
+        lock: RedisLock,
         tenant_id: str | None,
-    ) -> int | None:
+    ) -> tuple[int, int] | None:
         last_lock_time = time.monotonic()
 
         async_results = []
@@ -84,7 +85,7 @@ class RedisDocumentSet(RedisObjectHelper):
 
             async_results.append(result)
 
-        return len(async_results)
+        return len(async_results), len(async_results)
 
     def reset(self) -> None:
         self.redis.delete(self.taskset_key)

backend/danswer/redis/redis_object_helper.py

@@ -1,9 +1,9 @@
 from abc import ABC
 from abc import abstractmethod
 
-import redis
 from celery import Celery
 from redis import Redis
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.redis.redis_pool import get_redis_client
@@ -85,7 +85,13 @@ class RedisObjectHelper(ABC):
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
-        lock: redis.lock.Lock,
+        lock: RedisLock,
         tenant_id: str | None,
-    ) -> int | None:
-        pass
+    ) -> tuple[int, int] | None:
+        """First element should be the number of actual tasks generated, second should
+        be the number of docs that were candidates to be synced for the cc pair.
+
+        The need for this is when we are syncing stale docs referenced by multiple
+        connectors. In a single pass across multiple cc pairs, we only want a task
+        for be created for a particular document id the first time we see it.
+        The rest can be skipped."""

backend/danswer/redis/redis_usergroup.py

@@ -5,6 +5,7 @@ from uuid import uuid4
 import redis
 from celery import Celery
 from redis import Redis
+from redis.lock import Lock as RedisLock
 from sqlalchemy.orm import Session
 
 from danswer.configs.constants import CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT
@@ -51,15 +52,15 @@ class RedisUserGroup(RedisObjectHelper):
         celery_app: Celery,
         db_session: Session,
         redis_client: Redis,
-        lock: redis.lock.Lock,
+        lock: RedisLock,
         tenant_id: str | None,
-    ) -> int | None:
+    ) -> tuple[int, int] | None:
         last_lock_time = time.monotonic()
 
         async_results = []
 
         if not global_version.is_ee_version():
-            return 0
+            return 0, 0
 
         try:
             construct_document_select_by_usergroup = fetch_versioned_implementation(
@@ -67,7 +68,7 @@ class RedisUserGroup(RedisObjectHelper):
                 "construct_document_select_by_usergroup",
             )
         except ModuleNotFoundError:
-            return 0
+            return 0, 0
 
         stmt = construct_document_select_by_usergroup(int(self._id))
         for doc in db_session.scalars(stmt).yield_per(1):
@@ -97,7 +98,7 @@ class RedisUserGroup(RedisObjectHelper):
 
             async_results.append(result)
 
-        return len(async_results)
+        return len(async_results), len(async_results)
 
     def reset(self) -> None:
         self.redis.delete(self.taskset_key)

backend/danswer/server/manage/users.py

@@ -11,7 +11,6 @@ from fastapi import Body
 from fastapi import Depends
 from fastapi import HTTPException
 from fastapi import Request
-from fastapi import status
 from psycopg2.errors import UniqueViolation
 from pydantic import BaseModel
 from sqlalchemy import Column
@@ -27,6 +26,7 @@ from danswer.auth.noauth_user import fetch_no_auth_user
 from danswer.auth.noauth_user import set_no_auth_user_preferences
 from danswer.auth.schemas import UserRole
 from danswer.auth.schemas import UserStatus
+from danswer.auth.users import BasicAuthenticationError
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_curator_or_admin_user
 from danswer.auth.users import current_user
@@ -492,13 +492,10 @@ def verify_user_logged_in(
             store = get_kv_store()
             return fetch_no_auth_user(store)
 
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN, detail="User Not Authenticated"
-        )
+        raise BasicAuthenticationError(detail="User Not Authenticated")
 
     if user.oidc_expiry and user.oidc_expiry < datetime.now(timezone.utc):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
+        raise BasicAuthenticationError(
             detail="Access denied. User's OIDC token has expired.",
         )
 

backend/danswer/server/openai_assistants_api/asssistants_api.py

@@ -0,0 +1,273 @@
+from typing import Any
+from typing import Optional
+from uuid import uuid4
+
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi import Query
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from danswer.auth.users import current_user
+from danswer.db.engine import get_session
+from danswer.db.models import Persona
+from danswer.db.models import User
+from danswer.db.persona import get_persona_by_id
+from danswer.db.persona import get_personas
+from danswer.db.persona import mark_persona_as_deleted
+from danswer.db.persona import upsert_persona
+from danswer.db.persona import upsert_prompt
+from danswer.db.tools import get_tool_by_name
+from danswer.search.enums import RecencyBiasSetting
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+router = APIRouter(prefix="/assistants")
+
+
+# Base models
+class AssistantObject(BaseModel):
+    id: int
+    object: str = "assistant"
+    created_at: int
+    name: Optional[str] = None
+    description: Optional[str] = None
+    model: str
+    instructions: Optional[str] = None
+    tools: list[dict[str, Any]]
+    file_ids: list[str]
+    metadata: Optional[dict[str, Any]] = None
+
+
+class CreateAssistantRequest(BaseModel):
+    model: str
+    name: Optional[str] = None
+    description: Optional[str] = None
+    instructions: Optional[str] = None
+    tools: Optional[list[dict[str, Any]]] = None
+    file_ids: Optional[list[str]] = None
+    metadata: Optional[dict[str, Any]] = None
+
+
+class ModifyAssistantRequest(BaseModel):
+    model: Optional[str] = None
+    name: Optional[str] = None
+    description: Optional[str] = None
+    instructions: Optional[str] = None
+    tools: Optional[list[dict[str, Any]]] = None
+    file_ids: Optional[list[str]] = None
+    metadata: Optional[dict[str, Any]] = None
+
+
+class DeleteAssistantResponse(BaseModel):
+    id: int
+    object: str = "assistant.deleted"
+    deleted: bool
+
+
+class ListAssistantsResponse(BaseModel):
+    object: str = "list"
+    data: list[AssistantObject]
+    first_id: Optional[int] = None
+    last_id: Optional[int] = None
+    has_more: bool
+
+
+def persona_to_assistant(persona: Persona) -> AssistantObject:
+    return AssistantObject(
+        id=persona.id,
+        created_at=0,
+        name=persona.name,
+        description=persona.description,
+        model=persona.llm_model_version_override or "gpt-3.5-turbo",
+        instructions=persona.prompts[0].system_prompt if persona.prompts else None,
+        tools=[
+            {
+                "type": tool.display_name,
+                "function": {
+                    "name": tool.name,
+                    "description": tool.description,
+                    "schema": tool.openapi_schema,
+                },
+            }
+            for tool in persona.tools
+        ],
+        file_ids=[],  # Assuming no file support for now
+        metadata={},  # Assuming no metadata for now
+    )
+
+
+# API endpoints
+@router.post("")
+def create_assistant(
+    request: CreateAssistantRequest,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> AssistantObject:
+    prompt = None
+    if request.instructions:
+        prompt = upsert_prompt(
+            user=user,
+            name=f"Prompt for {request.name or 'New Assistant'}",
+            description="Auto-generated prompt",
+            system_prompt=request.instructions,
+            task_prompt="",
+            include_citations=True,
+            datetime_aware=True,
+            personas=[],
+            db_session=db_session,
+        )
+
+    tool_ids = []
+    for tool in request.tools or []:
+        tool_type = tool.get("type")
+        if not tool_type:
+            continue
+
+        try:
+            tool_db = get_tool_by_name(tool_type, db_session)
+            tool_ids.append(tool_db.id)
+        except ValueError:
+            # Skip tools that don't exist in the database
+            logger.error(f"Tool {tool_type} not found in database")
+            raise HTTPException(
+                status_code=404, detail=f"Tool {tool_type} not found in database"
+            )
+
+    persona = upsert_persona(
+        user=user,
+        name=request.name or f"Assistant-{uuid4()}",
+        description=request.description or "",
+        num_chunks=25,
+        llm_relevance_filter=True,
+        llm_filter_extraction=True,
+        recency_bias=RecencyBiasSetting.AUTO,
+        llm_model_provider_override=None,
+        llm_model_version_override=request.model,
+        starter_messages=None,
+        is_public=False,
+        db_session=db_session,
+        prompt_ids=[prompt.id] if prompt else [0],
+        document_set_ids=[],
+        tool_ids=tool_ids,
+        icon_color=None,
+        icon_shape=None,
+        is_visible=True,
+    )
+
+    if prompt:
+        prompt.personas = [persona]
+        db_session.commit()
+
+    return persona_to_assistant(persona)
+
+
+""
+
+
+@router.get("/{assistant_id}")
+def retrieve_assistant(
+    assistant_id: int,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> AssistantObject:
+    try:
+        persona = get_persona_by_id(
+            persona_id=assistant_id,
+            user=user,
+            db_session=db_session,
+            is_for_edit=False,
+        )
+    except ValueError:
+        persona = None
+
+    if not persona:
+        raise HTTPException(status_code=404, detail="Assistant not found")
+    return persona_to_assistant(persona)
+
+
+@router.post("/{assistant_id}")
+def modify_assistant(
+    assistant_id: int,
+    request: ModifyAssistantRequest,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> AssistantObject:
+    persona = get_persona_by_id(
+        persona_id=assistant_id,
+        user=user,
+        db_session=db_session,
+        is_for_edit=True,
+    )
+    if not persona:
+        raise HTTPException(status_code=404, detail="Assistant not found")
+
+    update_data = request.model_dump(exclude_unset=True)
+    for key, value in update_data.items():
+        setattr(persona, key, value)
+
+    if "instructions" in update_data and persona.prompts:
+        persona.prompts[0].system_prompt = update_data["instructions"]
+
+    db_session.commit()
+    return persona_to_assistant(persona)
+
+
+@router.delete("/{assistant_id}")
+def delete_assistant(
+    assistant_id: int,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> DeleteAssistantResponse:
+    try:
+        mark_persona_as_deleted(
+            persona_id=int(assistant_id),
+            user=user,
+            db_session=db_session,
+        )
+        return DeleteAssistantResponse(id=assistant_id, deleted=True)
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Assistant not found")
+
+
+@router.get("")
+def list_assistants(
+    limit: int = Query(20, le=100),
+    order: str = Query("desc", regex="^(asc|desc)$"),
+    after: Optional[int] = None,
+    before: Optional[int] = None,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> ListAssistantsResponse:
+    personas = list(
+        get_personas(
+            user=user,
+            db_session=db_session,
+            get_editable=False,
+            joinedload_all=True,
+        )
+    )
+
+    # Apply filtering based on after and before
+    if after:
+        personas = [p for p in personas if p.id > int(after)]
+    if before:
+        personas = [p for p in personas if p.id < int(before)]
+
+    # Apply ordering
+    personas.sort(key=lambda p: p.id, reverse=(order == "desc"))
+
+    # Apply limit
+    personas = personas[:limit]
+
+    assistants = [persona_to_assistant(p) for p in personas]
+
+    return ListAssistantsResponse(
+        data=assistants,
+        first_id=assistants[0].id if assistants else None,
+        last_id=assistants[-1].id if assistants else None,
+        has_more=len(personas) == limit,
+    )

backend/danswer/server/openai_assistants_api/full_openai_assistants_api.py

@@ -0,0 +1,19 @@
+from fastapi import APIRouter
+
+from danswer.server.openai_assistants_api.asssistants_api import (
+    router as assistants_router,
+)
+from danswer.server.openai_assistants_api.messages_api import router as messages_router
+from danswer.server.openai_assistants_api.runs_api import router as runs_router
+from danswer.server.openai_assistants_api.threads_api import router as threads_router
+
+
+def get_full_openai_assistants_api_router() -> APIRouter:
+    router = APIRouter(prefix="/openai-assistants")
+
+    router.include_router(assistants_router)
+    router.include_router(runs_router)
+    router.include_router(threads_router)
+    router.include_router(messages_router)
+
+    return router

backend/danswer/server/openai_assistants_api/messages_api.py

@@ -0,0 +1,235 @@
+import uuid
+from datetime import datetime
+from typing import Any
+from typing import Literal
+from typing import Optional
+
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import HTTPException
+from pydantic import BaseModel
+from pydantic import Field
+from sqlalchemy.orm import Session
+
+from danswer.auth.users import current_user
+from danswer.configs.constants import MessageType
+from danswer.db.chat import create_new_chat_message
+from danswer.db.chat import get_chat_message
+from danswer.db.chat import get_chat_messages_by_session
+from danswer.db.chat import get_chat_session_by_id
+from danswer.db.chat import get_or_create_root_message
+from danswer.db.engine import get_session
+from danswer.db.models import User
+from danswer.llm.utils import check_number_of_tokens
+
+router = APIRouter(prefix="")
+
+
+Role = Literal["user", "assistant"]
+
+
+class MessageContent(BaseModel):
+    type: Literal["text"]
+    text: str
+
+
+class Message(BaseModel):
+    id: str = Field(default_factory=lambda: f"msg_{uuid.uuid4()}")
+    object: Literal["thread.message"] = "thread.message"
+    created_at: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
+    thread_id: str
+    role: Role
+    content: list[MessageContent]
+    file_ids: list[str] = []
+    assistant_id: Optional[str] = None
+    run_id: Optional[str] = None
+    metadata: Optional[dict[str, Any]] = None  # Change this line to use dict[str, Any]
+
+
+class CreateMessageRequest(BaseModel):
+    role: Role
+    content: str
+    file_ids: list[str] = []
+    metadata: Optional[dict] = None
+
+
+class ListMessagesResponse(BaseModel):
+    object: Literal["list"] = "list"
+    data: list[Message]
+    first_id: str
+    last_id: str
+    has_more: bool
+
+
+@router.post("/threads/{thread_id}/messages")
+def create_message(
+    thread_id: str,
+    message: CreateMessageRequest,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> Message:
+    user_id = user.id if user else None
+
+    try:
+        chat_session = get_chat_session_by_id(
+            chat_session_id=uuid.UUID(thread_id),
+            user_id=user_id,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Chat session not found")
+
+    chat_messages = get_chat_messages_by_session(
+        chat_session_id=chat_session.id,
+        user_id=user.id if user else None,
+        db_session=db_session,
+    )
+    latest_message = (
+        chat_messages[-1]
+        if chat_messages
+        else get_or_create_root_message(chat_session.id, db_session)
+    )
+
+    new_message = create_new_chat_message(
+        chat_session_id=chat_session.id,
+        parent_message=latest_message,
+        message=message.content,
+        prompt_id=chat_session.persona.prompts[0].id,
+        token_count=check_number_of_tokens(message.content),
+        message_type=(
+            MessageType.USER if message.role == "user" else MessageType.ASSISTANT
+        ),
+        db_session=db_session,
+    )
+
+    return Message(
+        id=str(new_message.id),
+        thread_id=thread_id,
+        role="user",
+        content=[MessageContent(type="text", text=message.content)],
+        file_ids=message.file_ids,
+        metadata=message.metadata,
+    )
+
+
+@router.get("/threads/{thread_id}/messages")
+def list_messages(
+    thread_id: str,
+    limit: int = 20,
+    order: Literal["asc", "desc"] = "desc",
+    after: Optional[str] = None,
+    before: Optional[str] = None,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> ListMessagesResponse:
+    user_id = user.id if user else None
+
+    try:
+        chat_session = get_chat_session_by_id(
+            chat_session_id=uuid.UUID(thread_id),
+            user_id=user_id,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Chat session not found")
+
+    messages = get_chat_messages_by_session(
+        chat_session_id=chat_session.id,
+        user_id=user_id,
+        db_session=db_session,
+    )
+
+    # Apply filtering based on after and before
+    if after:
+        messages = [m for m in messages if str(m.id) >= after]
+    if before:
+        messages = [m for m in messages if str(m.id) <= before]
+
+    # Apply ordering
+    messages = sorted(messages, key=lambda m: m.id, reverse=(order == "desc"))
+
+    # Apply limit
+    messages = messages[:limit]
+
+    data = [
+        Message(
+            id=str(m.id),
+            thread_id=thread_id,
+            role="user" if m.message_type == "user" else "assistant",
+            content=[MessageContent(type="text", text=m.message)],
+            created_at=int(m.time_sent.timestamp()),
+        )
+        for m in messages
+    ]
+
+    return ListMessagesResponse(
+        data=data,
+        first_id=str(data[0].id) if data else "",
+        last_id=str(data[-1].id) if data else "",
+        has_more=len(messages) == limit,
+    )
+
+
+@router.get("/threads/{thread_id}/messages/{message_id}")
+def retrieve_message(
+    thread_id: str,
+    message_id: int,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> Message:
+    user_id = user.id if user else None
+
+    try:
+        chat_message = get_chat_message(
+            chat_message_id=message_id,
+            user_id=user_id,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Message not found")
+
+    return Message(
+        id=str(chat_message.id),
+        thread_id=thread_id,
+        role="user" if chat_message.message_type == "user" else "assistant",
+        content=[MessageContent(type="text", text=chat_message.message)],
+        created_at=int(chat_message.time_sent.timestamp()),
+    )
+
+
+class ModifyMessageRequest(BaseModel):
+    metadata: dict
+
+
+@router.post("/threads/{thread_id}/messages/{message_id}")
+def modify_message(
+    thread_id: str,
+    message_id: int,
+    request: ModifyMessageRequest,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> Message:
+    user_id = user.id if user else None
+
+    try:
+        chat_message = get_chat_message(
+            chat_message_id=message_id,
+            user_id=user_id,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Message not found")
+
+    # Update metadata
+    # TODO: Uncomment this once we have metadata in the chat message
+    # chat_message.metadata = request.metadata
+    # db_session.commit()
+
+    return Message(
+        id=str(chat_message.id),
+        thread_id=thread_id,
+        role="user" if chat_message.message_type == "user" else "assistant",
+        content=[MessageContent(type="text", text=chat_message.message)],
+        created_at=int(chat_message.time_sent.timestamp()),
+        metadata=request.metadata,
+    )

backend/danswer/server/openai_assistants_api/runs_api.py

@@ -0,0 +1,344 @@
+from typing import Literal
+from typing import Optional
+from uuid import UUID
+
+from fastapi import APIRouter
+from fastapi import BackgroundTasks
+from fastapi import Depends
+from fastapi import HTTPException
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from danswer.auth.users import current_user
+from danswer.chat.process_message import stream_chat_message_objects
+from danswer.configs.constants import MessageType
+from danswer.db.chat import create_new_chat_message
+from danswer.db.chat import get_chat_message
+from danswer.db.chat import get_chat_messages_by_session
+from danswer.db.chat import get_chat_session_by_id
+from danswer.db.chat import get_or_create_root_message
+from danswer.db.engine import get_session
+from danswer.db.models import ChatMessage
+from danswer.db.models import User
+from danswer.search.models import RetrievalDetails
+from danswer.server.query_and_chat.models import ChatMessageDetail
+from danswer.server.query_and_chat.models import CreateChatMessageRequest
+from danswer.tools.tool_implementations.search.search_tool import SearchTool
+from danswer.utils.logger import setup_logger
+
+
+logger = setup_logger()
+
+
+router = APIRouter()
+
+
+class RunRequest(BaseModel):
+    assistant_id: int
+    model: Optional[str] = None
+    instructions: Optional[str] = None
+    additional_instructions: Optional[str] = None
+    tools: Optional[list[dict]] = None
+    metadata: Optional[dict] = None
+
+
+RunStatus = Literal[
+    "queued",
+    "in_progress",
+    "requires_action",
+    "cancelling",
+    "cancelled",
+    "failed",
+    "completed",
+    "expired",
+]
+
+
+class RunResponse(BaseModel):
+    id: str
+    object: Literal["thread.run"]
+    created_at: int
+    assistant_id: int
+    thread_id: UUID
+    status: RunStatus
+    started_at: Optional[int] = None
+    expires_at: Optional[int] = None
+    cancelled_at: Optional[int] = None
+    failed_at: Optional[int] = None
+    completed_at: Optional[int] = None
+    last_error: Optional[dict] = None
+    model: str
+    instructions: str
+    tools: list[dict]
+    file_ids: list[str]
+    metadata: Optional[dict] = None
+
+
+def process_run_in_background(
+    message_id: int,
+    parent_message_id: int,
+    chat_session_id: UUID,
+    assistant_id: int,
+    instructions: str,
+    tools: list[dict],
+    user: User | None,
+    db_session: Session,
+) -> None:
+    # Get the latest message in the chat session
+    chat_session = get_chat_session_by_id(
+        chat_session_id=chat_session_id,
+        user_id=user.id if user else None,
+        db_session=db_session,
+    )
+
+    search_tool_retrieval_details = RetrievalDetails()
+    for tool in tools:
+        if tool["type"] == SearchTool.__name__ and (
+            retrieval_details := tool.get("retrieval_details")
+        ):
+            search_tool_retrieval_details = RetrievalDetails.model_validate(
+                retrieval_details
+            )
+            break
+
+    new_msg_req = CreateChatMessageRequest(
+        chat_session_id=chat_session_id,
+        parent_message_id=int(parent_message_id) if parent_message_id else None,
+        message=instructions,
+        file_descriptors=[],
+        prompt_id=chat_session.persona.prompts[0].id,
+        search_doc_ids=None,
+        retrieval_options=search_tool_retrieval_details,  # Adjust as needed
+        query_override=None,
+        regenerate=None,
+        llm_override=None,
+        prompt_override=None,
+        alternate_assistant_id=assistant_id,
+        use_existing_user_message=True,
+        existing_assistant_message_id=message_id,
+    )
+
+    run_message = get_chat_message(message_id, user.id if user else None, db_session)
+    try:
+        for packet in stream_chat_message_objects(
+            new_msg_req=new_msg_req,
+            user=user,
+            db_session=db_session,
+        ):
+            if isinstance(packet, ChatMessageDetail):
+                # Update the run status and message content
+                run_message = get_chat_message(
+                    message_id, user.id if user else None, db_session
+                )
+                if run_message:
+                    # this handles cancelling
+                    if run_message.error:
+                        return
+
+                    run_message.message = packet.message
+                    run_message.message_type = MessageType.ASSISTANT
+                    db_session.commit()
+    except Exception as e:
+        logger.exception("Error processing run in background")
+        run_message.error = str(e)
+        db_session.commit()
+        return
+
+    db_session.refresh(run_message)
+    if run_message.token_count == 0:
+        run_message.error = "No tokens generated"
+        db_session.commit()
+
+
+@router.post("/threads/{thread_id}/runs")
+def create_run(
+    thread_id: UUID,
+    run_request: RunRequest,
+    background_tasks: BackgroundTasks,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> RunResponse:
+    try:
+        chat_session = get_chat_session_by_id(
+            chat_session_id=thread_id,
+            user_id=user.id if user else None,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Thread not found")
+
+    chat_messages = get_chat_messages_by_session(
+        chat_session_id=chat_session.id,
+        user_id=user.id if user else None,
+        db_session=db_session,
+    )
+    latest_message = (
+        chat_messages[-1]
+        if chat_messages
+        else get_or_create_root_message(chat_session.id, db_session)
+    )
+
+    # Create a new "run" (chat message) in the session
+    new_message = create_new_chat_message(
+        chat_session_id=chat_session.id,
+        parent_message=latest_message,
+        message="",
+        prompt_id=chat_session.persona.prompts[0].id,
+        token_count=0,
+        message_type=MessageType.ASSISTANT,
+        db_session=db_session,
+        commit=False,
+    )
+    db_session.flush()
+    latest_message.latest_child_message = new_message.id
+    db_session.commit()
+
+    # Schedule the background task
+    background_tasks.add_task(
+        process_run_in_background,
+        new_message.id,
+        latest_message.id,
+        chat_session.id,
+        run_request.assistant_id,
+        run_request.instructions or "",
+        run_request.tools or [],
+        user,
+        db_session,
+    )
+
+    return RunResponse(
+        id=str(new_message.id),
+        object="thread.run",
+        created_at=int(new_message.time_sent.timestamp()),
+        assistant_id=run_request.assistant_id,
+        thread_id=chat_session.id,
+        status="queued",
+        model=run_request.model or "default_model",
+        instructions=run_request.instructions or "",
+        tools=run_request.tools or [],
+        file_ids=[],
+        metadata=run_request.metadata,
+    )
+
+
+@router.get("/threads/{thread_id}/runs/{run_id}")
+def retrieve_run(
+    thread_id: UUID,
+    run_id: str,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> RunResponse:
+    # Retrieve the chat message (which represents a "run" in DAnswer)
+    chat_message = get_chat_message(
+        chat_message_id=int(run_id),  # Convert string run_id to int
+        user_id=user.id if user else None,
+        db_session=db_session,
+    )
+    if not chat_message:
+        raise HTTPException(status_code=404, detail="Run not found")
+
+    chat_session = chat_message.chat_session
+
+    # Map DAnswer status to OpenAI status
+    run_status: RunStatus = "queued"
+    if chat_message.message:
+        run_status = "in_progress"
+    if chat_message.token_count != 0:
+        run_status = "completed"
+    if chat_message.error:
+        run_status = "cancelled"
+
+    return RunResponse(
+        id=run_id,
+        object="thread.run",
+        created_at=int(chat_message.time_sent.timestamp()),
+        assistant_id=chat_session.persona_id or 0,
+        thread_id=chat_session.id,
+        status=run_status,
+        started_at=int(chat_message.time_sent.timestamp()),
+        completed_at=(
+            int(chat_message.time_sent.timestamp()) if chat_message.message else None
+        ),
+        model=chat_session.current_alternate_model or "default_model",
+        instructions="",  # DAnswer doesn't store per-message instructions
+        tools=[],  # DAnswer doesn't have a direct equivalent for tools
+        file_ids=(
+            [file["id"] for file in chat_message.files] if chat_message.files else []
+        ),
+        metadata=None,  # DAnswer doesn't store metadata for individual messages
+    )
+
+
+@router.post("/threads/{thread_id}/runs/{run_id}/cancel")
+def cancel_run(
+    thread_id: UUID,
+    run_id: str,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> RunResponse:
+    # In DAnswer, we don't have a direct equivalent to cancelling a run
+    # We'll simulate it by marking the message as "cancelled"
+    chat_message = (
+        db_session.query(ChatMessage).filter(ChatMessage.id == run_id).first()
+    )
+    if not chat_message:
+        raise HTTPException(status_code=404, detail="Run not found")
+
+    chat_message.error = "Cancelled"
+    db_session.commit()
+
+    return retrieve_run(thread_id, run_id, user, db_session)
+
+
+@router.get("/threads/{thread_id}/runs")
+def list_runs(
+    thread_id: UUID,
+    limit: int = 20,
+    order: Literal["asc", "desc"] = "desc",
+    after: Optional[str] = None,
+    before: Optional[str] = None,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> list[RunResponse]:
+    # In DAnswer, we'll treat each message in a chat session as a "run"
+    chat_messages = get_chat_messages_by_session(
+        chat_session_id=thread_id,
+        user_id=user.id if user else None,
+        db_session=db_session,
+    )
+
+    # Apply pagination
+    if after:
+        chat_messages = [msg for msg in chat_messages if str(msg.id) > after]
+    if before:
+        chat_messages = [msg for msg in chat_messages if str(msg.id) < before]
+
+    # Apply ordering
+    chat_messages = sorted(
+        chat_messages, key=lambda msg: msg.time_sent, reverse=(order == "desc")
+    )
+
+    # Apply limit
+    chat_messages = chat_messages[:limit]
+
+    return [
+        retrieve_run(thread_id, str(msg.id), user, db_session) for msg in chat_messages
+    ]
+
+
+@router.get("/threads/{thread_id}/runs/{run_id}/steps")
+def list_run_steps(
+    run_id: str,
+    limit: int = 20,
+    order: Literal["asc", "desc"] = "desc",
+    after: Optional[str] = None,
+    before: Optional[str] = None,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> list[dict]:  # You may want to create a specific model for run steps
+    # DAnswer doesn't have an equivalent to run steps
+    # We'll return an empty list to maintain API compatibility
+    return []
+
+
+# Additional helper functions can be added here if needed

backend/danswer/server/openai_assistants_api/threads_api.py

@@ -0,0 +1,156 @@
+from typing import Optional
+from uuid import UUID
+
+from fastapi import APIRouter
+from fastapi import Depends
+from fastapi import HTTPException
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from danswer.auth.users import current_user
+from danswer.db.chat import create_chat_session
+from danswer.db.chat import delete_chat_session
+from danswer.db.chat import get_chat_session_by_id
+from danswer.db.chat import get_chat_sessions_by_user
+from danswer.db.chat import update_chat_session
+from danswer.db.engine import get_session
+from danswer.db.models import User
+from danswer.server.query_and_chat.models import ChatSessionDetails
+from danswer.server.query_and_chat.models import ChatSessionsResponse
+
+router = APIRouter(prefix="/threads")
+
+
+# Models
+class Thread(BaseModel):
+    id: UUID
+    object: str = "thread"
+    created_at: int
+    metadata: Optional[dict[str, str]] = None
+
+
+class CreateThreadRequest(BaseModel):
+    messages: Optional[list[dict]] = None
+    metadata: Optional[dict[str, str]] = None
+
+
+class ModifyThreadRequest(BaseModel):
+    metadata: Optional[dict[str, str]] = None
+
+
+# API Endpoints
+@router.post("")
+def create_thread(
+    request: CreateThreadRequest,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> Thread:
+    user_id = user.id if user else None
+    new_chat_session = create_chat_session(
+        db_session=db_session,
+        description="",  # Leave the naming till later to prevent delay
+        user_id=user_id,
+        persona_id=0,
+    )
+
+    return Thread(
+        id=new_chat_session.id,
+        created_at=int(new_chat_session.time_created.timestamp()),
+        metadata=request.metadata,
+    )
+
+
+@router.get("/{thread_id}")
+def retrieve_thread(
+    thread_id: UUID,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> Thread:
+    user_id = user.id if user else None
+    try:
+        chat_session = get_chat_session_by_id(
+            chat_session_id=thread_id,
+            user_id=user_id,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Thread not found")
+
+    return Thread(
+        id=chat_session.id,
+        created_at=int(chat_session.time_created.timestamp()),
+        metadata=None,  # Assuming we don't store metadata in our current implementation
+    )
+
+
+@router.post("/{thread_id}")
+def modify_thread(
+    thread_id: UUID,
+    request: ModifyThreadRequest,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> Thread:
+    user_id = user.id if user else None
+    try:
+        chat_session = update_chat_session(
+            db_session=db_session,
+            user_id=user_id,
+            chat_session_id=thread_id,
+            description=None,  # Not updating description
+            sharing_status=None,  # Not updating sharing status
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Thread not found")
+
+    return Thread(
+        id=chat_session.id,
+        created_at=int(chat_session.time_created.timestamp()),
+        metadata=request.metadata,
+    )
+
+
+@router.delete("/{thread_id}")
+def delete_thread(
+    thread_id: UUID,
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> dict:
+    user_id = user.id if user else None
+    try:
+        delete_chat_session(
+            user_id=user_id,
+            chat_session_id=thread_id,
+            db_session=db_session,
+        )
+    except ValueError:
+        raise HTTPException(status_code=404, detail="Thread not found")
+
+    return {"id": str(thread_id), "object": "thread.deleted", "deleted": True}
+
+
+@router.get("")
+def list_threads(
+    user: User | None = Depends(current_user),
+    db_session: Session = Depends(get_session),
+) -> ChatSessionsResponse:
+    user_id = user.id if user else None
+    chat_sessions = get_chat_sessions_by_user(
+        user_id=user_id,
+        deleted=False,
+        db_session=db_session,
+    )
+
+    return ChatSessionsResponse(
+        sessions=[
+            ChatSessionDetails(
+                id=chat.id,
+                name=chat.description,
+                persona_id=chat.persona_id,
+                time_created=chat.time_created.isoformat(),
+                shared_status=chat.shared_status,
+                folder_id=chat.folder_id,
+                current_alternate_model=chat.current_alternate_model,
+            )
+            for chat in chat_sessions
+        ]
+    )

backend/danswer/server/query_and_chat/chat_backend.py

@@ -347,7 +347,6 @@ def handle_new_chat_message(
             for packet in stream_chat_message(
                 new_msg_req=chat_message_req,
                 user=user,
-                use_existing_user_message=chat_message_req.use_existing_user_message,
                 litellm_additional_headers=extract_headers(
                     request.headers, LITELLM_PASS_THROUGH_HEADERS
                 ),

backend/danswer/server/query_and_chat/models.py

@@ -108,6 +108,9 @@ class CreateChatMessageRequest(ChunkContext):
     # used for seeded chats to kick off the generation of an AI answer
     use_existing_user_message: bool = False
 
+    # used for "OpenAI Assistants API"
+    existing_assistant_message_id: int | None = None
+
     # forces the LLM to return a structured response, see
     # https://platform.openai.com/docs/guides/structured-outputs/introduction
     structured_response_format: dict | None = None

backend/danswer/tools/tool_constructor.py

@@ -0,0 +1,255 @@
+from typing import cast
+from uuid import UUID
+
+from pydantic import BaseModel
+from pydantic import Field
+from sqlalchemy.orm import Session
+
+from danswer.configs.app_configs import AZURE_DALLE_API_BASE
+from danswer.configs.app_configs import AZURE_DALLE_API_KEY
+from danswer.configs.app_configs import AZURE_DALLE_API_VERSION
+from danswer.configs.app_configs import AZURE_DALLE_DEPLOYMENT_NAME
+from danswer.configs.chat_configs import BING_API_KEY
+from danswer.configs.model_configs import GEN_AI_TEMPERATURE
+from danswer.db.llm import fetch_existing_llm_providers
+from danswer.db.models import Persona
+from danswer.db.models import User
+from danswer.file_store.models import InMemoryChatFile
+from danswer.llm.answering.models import AnswerStyleConfig
+from danswer.llm.answering.models import CitationConfig
+from danswer.llm.answering.models import DocumentPruningConfig
+from danswer.llm.answering.models import PromptConfig
+from danswer.llm.interfaces import LLM
+from danswer.llm.interfaces import LLMConfig
+from danswer.natural_language_processing.utils import get_tokenizer
+from danswer.search.enums import LLMEvaluationType
+from danswer.search.models import InferenceSection
+from danswer.search.models import RetrievalDetails
+from danswer.tools.built_in_tools import get_built_in_tool_by_id
+from danswer.tools.models import DynamicSchemaInfo
+from danswer.tools.tool import Tool
+from danswer.tools.tool_implementations.custom.custom_tool import (
+    build_custom_tools_from_openapi_schema_and_headers,
+)
+from danswer.tools.tool_implementations.images.image_generation_tool import (
+    ImageGenerationTool,
+)
+from danswer.tools.tool_implementations.internet_search.internet_search_tool import (
+    InternetSearchTool,
+)
+from danswer.tools.tool_implementations.search.search_tool import SearchTool
+from danswer.tools.utils import compute_all_tool_tokens
+from danswer.tools.utils import explicit_tool_calling_supported
+from danswer.utils.headers import header_dict_to_header_list
+from danswer.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def _get_image_generation_config(llm: LLM, db_session: Session) -> LLMConfig:
+    """Helper function to get image generation LLM config based on available providers"""
+    if llm and llm.config.api_key and llm.config.model_provider == "openai":
+        return LLMConfig(
+            model_provider=llm.config.model_provider,
+            model_name="dall-e-3",
+            temperature=GEN_AI_TEMPERATURE,
+            api_key=llm.config.api_key,
+            api_base=llm.config.api_base,
+            api_version=llm.config.api_version,
+        )
+
+    if llm.config.model_provider == "azure" and AZURE_DALLE_API_KEY is not None:
+        return LLMConfig(
+            model_provider="azure",
+            model_name=f"azure/{AZURE_DALLE_DEPLOYMENT_NAME}",
+            temperature=GEN_AI_TEMPERATURE,
+            api_key=AZURE_DALLE_API_KEY,
+            api_base=AZURE_DALLE_API_BASE,
+            api_version=AZURE_DALLE_API_VERSION,
+        )
+
+    # Fallback to checking for OpenAI provider in database
+    llm_providers = fetch_existing_llm_providers(db_session)
+    openai_provider = next(
+        iter(
+            [
+                llm_provider
+                for llm_provider in llm_providers
+                if llm_provider.provider == "openai"
+            ]
+        ),
+        None,
+    )
+
+    if not openai_provider or not openai_provider.api_key:
+        raise ValueError("Image generation tool requires an OpenAI API key")
+
+    return LLMConfig(
+        model_provider=openai_provider.provider,
+        model_name="dall-e-3",
+        temperature=GEN_AI_TEMPERATURE,
+        api_key=openai_provider.api_key,
+        api_base=openai_provider.api_base,
+        api_version=openai_provider.api_version,
+    )
+
+
+class SearchToolConfig(BaseModel):
+    answer_style_config: AnswerStyleConfig = Field(
+        default_factory=lambda: AnswerStyleConfig(citation_config=CitationConfig())
+    )
+    document_pruning_config: DocumentPruningConfig = Field(
+        default_factory=DocumentPruningConfig
+    )
+    retrieval_options: RetrievalDetails = Field(default_factory=RetrievalDetails)
+    selected_sections: list[InferenceSection] | None = None
+    chunks_above: int = 0
+    chunks_below: int = 0
+    full_doc: bool = False
+    latest_query_files: list[InMemoryChatFile] | None = None
+
+
+class InternetSearchToolConfig(BaseModel):
+    answer_style_config: AnswerStyleConfig = Field(
+        default_factory=lambda: AnswerStyleConfig(
+            citation_config=CitationConfig(all_docs_useful=True)
+        )
+    )
+
+
+class ImageGenerationToolConfig(BaseModel):
+    additional_headers: dict[str, str] | None = None
+
+
+class CustomToolConfig(BaseModel):
+    chat_session_id: UUID | None = None
+    message_id: int | None = None
+    additional_headers: dict[str, str] | None = None
+
+
+def construct_tools(
+    persona: Persona,
+    prompt_config: PromptConfig,
+    db_session: Session,
+    user: User | None,
+    llm: LLM,
+    fast_llm: LLM,
+    search_tool_config: SearchToolConfig | None = None,
+    internet_search_tool_config: InternetSearchToolConfig | None = None,
+    image_generation_tool_config: ImageGenerationToolConfig | None = None,
+    custom_tool_config: CustomToolConfig | None = None,
+) -> dict[int, list[Tool]]:
+    """Constructs tools based on persona configuration and available APIs"""
+    tool_dict: dict[int, list[Tool]] = {}
+
+    for db_tool_model in persona.tools:
+        if db_tool_model.in_code_tool_id:
+            tool_cls = get_built_in_tool_by_id(db_tool_model.id, db_session)
+
+            # Handle Search Tool
+            if tool_cls.__name__ == SearchTool.__name__:
+                if not search_tool_config:
+                    search_tool_config = SearchToolConfig()
+
+                search_tool = SearchTool(
+                    db_session=db_session,
+                    user=user,
+                    persona=persona,
+                    retrieval_options=search_tool_config.retrieval_options,
+                    prompt_config=prompt_config,
+                    llm=llm,
+                    fast_llm=fast_llm,
+                    pruning_config=search_tool_config.document_pruning_config,
+                    answer_style_config=search_tool_config.answer_style_config,
+                    selected_sections=search_tool_config.selected_sections,
+                    chunks_above=search_tool_config.chunks_above,
+                    chunks_below=search_tool_config.chunks_below,
+                    full_doc=search_tool_config.full_doc,
+                    evaluation_type=(
+                        LLMEvaluationType.BASIC
+                        if persona.llm_relevance_filter
+                        else LLMEvaluationType.SKIP
+                    ),
+                )
+                tool_dict[db_tool_model.id] = [search_tool]
+
+            # Handle Image Generation Tool
+            elif tool_cls.__name__ == ImageGenerationTool.__name__:
+                if not image_generation_tool_config:
+                    image_generation_tool_config = ImageGenerationToolConfig()
+
+                img_generation_llm_config = _get_image_generation_config(
+                    llm, db_session
+                )
+
+                tool_dict[db_tool_model.id] = [
+                    ImageGenerationTool(
+                        api_key=cast(str, img_generation_llm_config.api_key),
+                        api_base=img_generation_llm_config.api_base,
+                        api_version=img_generation_llm_config.api_version,
+                        additional_headers=image_generation_tool_config.additional_headers,
+                        model=img_generation_llm_config.model_name,
+                    )
+                ]
+
+            # Handle Internet Search Tool
+            elif tool_cls.__name__ == InternetSearchTool.__name__:
+                if not internet_search_tool_config:
+                    internet_search_tool_config = InternetSearchToolConfig()
+
+                if not BING_API_KEY:
+                    raise ValueError(
+                        "Internet search tool requires a Bing API key, please contact your Danswer admin to get it added!"
+                    )
+                tool_dict[db_tool_model.id] = [
+                    InternetSearchTool(
+                        api_key=BING_API_KEY,
+                        answer_style_config=internet_search_tool_config.answer_style_config,
+                        prompt_config=prompt_config,
+                    )
+                ]
+
+        # Handle custom tools
+        elif db_tool_model.openapi_schema:
+            if not custom_tool_config:
+                custom_tool_config = CustomToolConfig()
+
+            tool_dict[db_tool_model.id] = cast(
+                list[Tool],
+                build_custom_tools_from_openapi_schema_and_headers(
+                    db_tool_model.openapi_schema,
+                    dynamic_schema_info=DynamicSchemaInfo(
+                        chat_session_id=custom_tool_config.chat_session_id,
+                        message_id=custom_tool_config.message_id,
+                    ),
+                    custom_headers=(db_tool_model.custom_headers or [])
+                    + (
+                        header_dict_to_header_list(
+                            custom_tool_config.additional_headers or {}
+                        )
+                    ),
+                ),
+            )
+
+    tools: list[Tool] = []
+    for tool_list in tool_dict.values():
+        tools.extend(tool_list)
+
+    # factor in tool definition size when pruning
+    if search_tool_config:
+        search_tool_config.document_pruning_config.tool_num_tokens = (
+            compute_all_tool_tokens(
+                tools,
+                get_tokenizer(
+                    model_name=llm.config.model_name,
+                    provider_type=llm.config.model_provider,
+                ),
+            )
+        )
+        search_tool_config.document_pruning_config.using_tool_message = (
+            explicit_tool_calling_supported(
+                llm.config.model_provider, llm.config.model_name
+            )
+        )
+
+    return tool_dict

backend/danswer/utils/variable_functionality.py

@@ -1,5 +1,6 @@
 import functools
 import importlib
+import inspect
 from typing import Any
 from typing import TypeVar
 
@@ -139,8 +140,19 @@ def fetch_ee_implementation_or_noop(
         Exception: If EE is enabled but the fetch fails.
     """
     if not global_version.is_ee_version():
-        return lambda *args, **kwargs: noop_return_value
+        if inspect.iscoroutinefunction(noop_return_value):
 
+            async def async_noop(*args: Any, **kwargs: Any) -> Any:
+                return await noop_return_value(*args, **kwargs)
+
+            return async_noop
+
+        else:
+
+            def sync_noop(*args: Any, **kwargs: Any) -> Any:
+                return noop_return_value
+
+            return sync_noop
     try:
         return fetch_versioned_implementation(module, attribute)
     except Exception as e:

backend/tests/integration/common_utils/managers/user.py

@@ -13,6 +13,14 @@ from tests.integration.common_utils.constants import GENERAL_HEADERS
 from tests.integration.common_utils.test_models import DATestUser
 
 
+DOMAIN = "test.com"
+DEFAULT_PASSWORD = "test"
+
+
+def build_email(name: str) -> str:
+    return f"{name}@test.com"
+
+
 class UserManager:
     @staticmethod
     def create(
@@ -23,9 +31,9 @@ class UserManager:
             name = f"test{str(uuid4())}"
 
         if email is None:
-            email = f"{name}@test.com"
+            email = build_email(name)
 
-        password = "test"
+        password = DEFAULT_PASSWORD
 
         body = {
             "email": email,

backend/tests/integration/openai_assistants_api/conftest.py

@@ -0,0 +1,55 @@
+from typing import Optional
+from uuid import UUID
+
+import pytest
+import requests
+
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.constants import GENERAL_HEADERS
+from tests.integration.common_utils.managers.llm_provider import LLMProviderManager
+from tests.integration.common_utils.managers.user import build_email
+from tests.integration.common_utils.managers.user import DEFAULT_PASSWORD
+from tests.integration.common_utils.managers.user import UserManager
+from tests.integration.common_utils.test_models import DATestLLMProvider
+from tests.integration.common_utils.test_models import DATestUser
+
+BASE_URL = f"{API_SERVER_URL}/openai-assistants"
+
+
+@pytest.fixture
+def admin_user() -> DATestUser | None:
+    try:
+        return UserManager.create("admin_user")
+    except Exception:
+        pass
+
+    try:
+        return UserManager.login_as_user(
+            DATestUser(
+                id="",
+                email=build_email("admin_user"),
+                password=DEFAULT_PASSWORD,
+                headers=GENERAL_HEADERS,
+            )
+        )
+    except Exception:
+        pass
+
+    return None
+
+
+@pytest.fixture
+def llm_provider(admin_user: DATestUser | None) -> DATestLLMProvider:
+    return LLMProviderManager.create(user_performing_action=admin_user)
+
+
+@pytest.fixture
+def thread_id(admin_user: Optional[DATestUser]) -> UUID:
+    # Create a thread to use in the tests
+    response = requests.post(
+        f"{BASE_URL}/threads",  # Updated endpoint path
+        json={},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    return UUID(response.json()["id"])

backend/tests/integration/openai_assistants_api/test_assistants.py

@@ -0,0 +1,151 @@
+import requests
+
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.constants import GENERAL_HEADERS
+from tests.integration.common_utils.test_models import DATestUser
+
+ASSISTANTS_URL = f"{API_SERVER_URL}/openai-assistants/assistants"
+
+
+def test_create_assistant(admin_user: DATestUser | None) -> None:
+    response = requests.post(
+        ASSISTANTS_URL,
+        json={
+            "model": "gpt-3.5-turbo",
+            "name": "Test Assistant",
+            "description": "A test assistant",
+            "instructions": "You are a helpful assistant.",
+        },
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["name"] == "Test Assistant"
+    assert data["description"] == "A test assistant"
+    assert data["model"] == "gpt-3.5-turbo"
+    assert data["instructions"] == "You are a helpful assistant."
+
+
+def test_retrieve_assistant(admin_user: DATestUser | None) -> None:
+    # First, create an assistant
+    create_response = requests.post(
+        ASSISTANTS_URL,
+        json={"model": "gpt-3.5-turbo", "name": "Retrieve Test"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert create_response.status_code == 200
+    assistant_id = create_response.json()["id"]
+
+    # Now, retrieve the assistant
+    response = requests.get(
+        f"{ASSISTANTS_URL}/{assistant_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["id"] == assistant_id
+    assert data["name"] == "Retrieve Test"
+
+
+def test_modify_assistant(admin_user: DATestUser | None) -> None:
+    # First, create an assistant
+    create_response = requests.post(
+        ASSISTANTS_URL,
+        json={"model": "gpt-3.5-turbo", "name": "Modify Test"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert create_response.status_code == 200
+    assistant_id = create_response.json()["id"]
+
+    # Now, modify the assistant
+    response = requests.post(
+        f"{ASSISTANTS_URL}/{assistant_id}",
+        json={"name": "Modified Assistant", "instructions": "New instructions"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["id"] == assistant_id
+    assert data["name"] == "Modified Assistant"
+    assert data["instructions"] == "New instructions"
+
+
+def test_delete_assistant(admin_user: DATestUser | None) -> None:
+    # First, create an assistant
+    create_response = requests.post(
+        ASSISTANTS_URL,
+        json={"model": "gpt-3.5-turbo", "name": "Delete Test"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert create_response.status_code == 200
+    assistant_id = create_response.json()["id"]
+
+    # Now, delete the assistant
+    response = requests.delete(
+        f"{ASSISTANTS_URL}/{assistant_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["id"] == assistant_id
+    assert data["deleted"] is True
+
+
+def test_list_assistants(admin_user: DATestUser | None) -> None:
+    # Create multiple assistants
+    for i in range(3):
+        requests.post(
+            ASSISTANTS_URL,
+            json={"model": "gpt-3.5-turbo", "name": f"List Test {i}"},
+            headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+        )
+
+    # Now, list the assistants
+    response = requests.get(
+        ASSISTANTS_URL,
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert data["object"] == "list"
+    assert len(data["data"]) >= 3  # At least the 3 we just created
+    assert all(assistant["object"] == "assistant" for assistant in data["data"])
+
+
+def test_list_assistants_pagination(admin_user: DATestUser | None) -> None:
+    # Create 5 assistants
+    for i in range(5):
+        requests.post(
+            ASSISTANTS_URL,
+            json={"model": "gpt-3.5-turbo", "name": f"Pagination Test {i}"},
+            headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+        )
+
+    # List assistants with limit
+    response = requests.get(
+        f"{ASSISTANTS_URL}?limit=2",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert len(data["data"]) == 2
+    assert data["has_more"] is True
+
+    # Get next page
+    before = data["last_id"]
+    response = requests.get(
+        f"{ASSISTANTS_URL}?limit=2&before={before}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert len(data["data"]) == 2
+
+
+def test_assistant_not_found(admin_user: DATestUser | None) -> None:
+    non_existent_id = -99
+    response = requests.get(
+        f"{ASSISTANTS_URL}/{non_existent_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 404

backend/tests/integration/openai_assistants_api/test_messages.py

@@ -0,0 +1,133 @@
+import uuid
+from typing import Optional
+
+import pytest
+import requests
+
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.constants import GENERAL_HEADERS
+from tests.integration.common_utils.test_models import DATestUser
+
+BASE_URL = f"{API_SERVER_URL}/openai-assistants/threads"
+
+
+@pytest.fixture
+def thread_id(admin_user: Optional[DATestUser]) -> str:
+    response = requests.post(
+        BASE_URL,
+        json={},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    return response.json()["id"]
+
+
+def test_create_message(admin_user: Optional[DATestUser], thread_id: str) -> None:
+    response = requests.post(
+        f"{BASE_URL}/{thread_id}/messages",  # URL structure matches API
+        json={
+            "role": "user",
+            "content": "Hello, world!",
+            "file_ids": [],
+            "metadata": {"key": "value"},
+        },
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "id" in response_json
+    assert response_json["thread_id"] == thread_id
+    assert response_json["role"] == "user"
+    assert response_json["content"] == [{"type": "text", "text": "Hello, world!"}]
+    assert response_json["metadata"] == {"key": "value"}
+
+
+def test_list_messages(admin_user: Optional[DATestUser], thread_id: str) -> None:
+    # Create a message first
+    requests.post(
+        f"{BASE_URL}/{thread_id}/messages",
+        json={"role": "user", "content": "Test message"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+
+    # Now, list the messages
+    response = requests.get(
+        f"{BASE_URL}/{thread_id}/messages",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["object"] == "list"
+    assert isinstance(response_json["data"], list)
+    assert len(response_json["data"]) > 0
+    assert "first_id" in response_json
+    assert "last_id" in response_json
+    assert "has_more" in response_json
+
+
+def test_retrieve_message(admin_user: Optional[DATestUser], thread_id: str) -> None:
+    # Create a message first
+    create_response = requests.post(
+        f"{BASE_URL}/{thread_id}/messages",
+        json={"role": "user", "content": "Test message"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    message_id = create_response.json()["id"]
+
+    # Now, retrieve the message
+    response = requests.get(
+        f"{BASE_URL}/{thread_id}/messages/{message_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["id"] == message_id
+    assert response_json["thread_id"] == thread_id
+    assert response_json["role"] == "user"
+    assert response_json["content"] == [{"type": "text", "text": "Test message"}]
+
+
+def test_modify_message(admin_user: Optional[DATestUser], thread_id: str) -> None:
+    # Create a message first
+    create_response = requests.post(
+        f"{BASE_URL}/{thread_id}/messages",
+        json={"role": "user", "content": "Test message"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    message_id = create_response.json()["id"]
+
+    # Now, modify the message
+    response = requests.post(
+        f"{BASE_URL}/{thread_id}/messages/{message_id}",
+        json={"metadata": {"new_key": "new_value"}},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert response_json["id"] == message_id
+    assert response_json["thread_id"] == thread_id
+    assert response_json["metadata"] == {"new_key": "new_value"}
+
+
+def test_error_handling(admin_user: Optional[DATestUser]) -> None:
+    non_existent_thread_id = str(uuid.uuid4())
+    non_existent_message_id = -99
+
+    # Test with non-existent thread
+    response = requests.post(
+        f"{BASE_URL}/{non_existent_thread_id}/messages",
+        json={"role": "user", "content": "Test message"},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 404
+
+    # Test with non-existent message
+    response = requests.get(
+        f"{BASE_URL}/{non_existent_thread_id}/messages/{non_existent_message_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 404

backend/tests/integration/openai_assistants_api/test_runs.py

@@ -0,0 +1,137 @@
+from uuid import UUID
+
+import pytest
+import requests
+
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.constants import GENERAL_HEADERS
+from tests.integration.common_utils.test_models import DATestLLMProvider
+from tests.integration.common_utils.test_models import DATestUser
+
+BASE_URL = f"{API_SERVER_URL}/openai-assistants"
+
+
+@pytest.fixture
+def run_id(admin_user: DATestUser | None, thread_id: UUID) -> str:
+    """Create a run and return its ID."""
+    response = requests.post(
+        f"{BASE_URL}/threads/{thread_id}/runs",
+        json={
+            "assistant_id": 0,
+        },
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+    return response.json()["id"]
+
+
+def test_create_run(
+    admin_user: DATestUser | None, thread_id: UUID, llm_provider: DATestLLMProvider
+) -> None:
+    response = requests.post(
+        f"{BASE_URL}/threads/{thread_id}/runs",
+        json={
+            "assistant_id": 0,
+            "model": "gpt-3.5-turbo",
+            "instructions": "Test instructions",
+        },
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "id" in response_json
+    assert response_json["object"] == "thread.run"
+    assert "created_at" in response_json
+    assert response_json["assistant_id"] == 0
+    assert UUID(response_json["thread_id"]) == thread_id
+    assert response_json["status"] == "queued"
+    assert response_json["model"] == "gpt-3.5-turbo"
+    assert response_json["instructions"] == "Test instructions"
+
+
+def test_retrieve_run(
+    admin_user: DATestUser | None,
+    thread_id: UUID,
+    run_id: str,
+    llm_provider: DATestLLMProvider,
+) -> None:
+    retrieve_response = requests.get(
+        f"{BASE_URL}/threads/{thread_id}/runs/{run_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert retrieve_response.status_code == 200
+
+    response_json = retrieve_response.json()
+    assert response_json["id"] == run_id
+    assert response_json["object"] == "thread.run"
+    assert "created_at" in response_json
+    assert UUID(response_json["thread_id"]) == thread_id
+
+
+def test_cancel_run(
+    admin_user: DATestUser | None,
+    thread_id: UUID,
+    run_id: str,
+    llm_provider: DATestLLMProvider,
+) -> None:
+    cancel_response = requests.post(
+        f"{BASE_URL}/threads/{thread_id}/runs/{run_id}/cancel",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert cancel_response.status_code == 200
+
+    response_json = cancel_response.json()
+    assert response_json["id"] == run_id
+    assert response_json["status"] == "cancelled"
+
+
+def test_list_runs(
+    admin_user: DATestUser | None, thread_id: UUID, llm_provider: DATestLLMProvider
+) -> None:
+    # Create a few runs
+    for _ in range(3):
+        requests.post(
+            f"{BASE_URL}/threads/{thread_id}/runs",
+            json={
+                "assistant_id": 0,
+            },
+            headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+        )
+
+    # Now, list the runs
+    list_response = requests.get(
+        f"{BASE_URL}/threads/{thread_id}/runs",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert list_response.status_code == 200
+
+    response_json = list_response.json()
+    assert isinstance(response_json, list)
+    assert len(response_json) >= 3
+
+    for run in response_json:
+        assert "id" in run
+        assert run["object"] == "thread.run"
+        assert "created_at" in run
+        assert UUID(run["thread_id"]) == thread_id
+        assert "status" in run
+        assert "model" in run
+
+
+def test_list_run_steps(
+    admin_user: DATestUser | None,
+    thread_id: UUID,
+    run_id: str,
+    llm_provider: DATestLLMProvider,
+) -> None:
+    steps_response = requests.get(
+        f"{BASE_URL}/threads/{thread_id}/runs/{run_id}/steps",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert steps_response.status_code == 200
+
+    response_json = steps_response.json()
+    assert isinstance(response_json, list)
+    # Since DAnswer doesn't have an equivalent to run steps, we expect an empty list
+    assert len(response_json) == 0

backend/tests/integration/openai_assistants_api/test_threads.py

@@ -0,0 +1,132 @@
+from uuid import UUID
+
+import requests
+
+from danswer.db.models import ChatSessionSharedStatus
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.constants import GENERAL_HEADERS
+from tests.integration.common_utils.test_models import DATestUser
+
+THREADS_URL = f"{API_SERVER_URL}/openai-assistants/threads"
+
+
+def test_create_thread(admin_user: DATestUser | None) -> None:
+    response = requests.post(
+        THREADS_URL,
+        json={"messages": None, "metadata": {"key": "value"}},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert response.status_code == 200
+
+    response_json = response.json()
+    assert "id" in response_json
+    assert response_json["object"] == "thread"
+    assert "created_at" in response_json
+    assert response_json["metadata"] == {"key": "value"}
+
+
+def test_retrieve_thread(admin_user: DATestUser | None) -> None:
+    # First, create a thread
+    create_response = requests.post(
+        THREADS_URL,
+        json={},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert create_response.status_code == 200
+    thread_id = create_response.json()["id"]
+
+    # Now, retrieve the thread
+    retrieve_response = requests.get(
+        f"{THREADS_URL}/{thread_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert retrieve_response.status_code == 200
+
+    response_json = retrieve_response.json()
+    assert response_json["id"] == thread_id
+    assert response_json["object"] == "thread"
+    assert "created_at" in response_json
+
+
+def test_modify_thread(admin_user: DATestUser | None) -> None:
+    # First, create a thread
+    create_response = requests.post(
+        THREADS_URL,
+        json={},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert create_response.status_code == 200
+    thread_id = create_response.json()["id"]
+
+    # Now, modify the thread
+    modify_response = requests.post(
+        f"{THREADS_URL}/{thread_id}",
+        json={"metadata": {"new_key": "new_value"}},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert modify_response.status_code == 200
+
+    response_json = modify_response.json()
+    assert response_json["id"] == thread_id
+    assert response_json["metadata"] == {"new_key": "new_value"}
+
+
+def test_delete_thread(admin_user: DATestUser | None) -> None:
+    # First, create a thread
+    create_response = requests.post(
+        THREADS_URL,
+        json={},
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert create_response.status_code == 200
+    thread_id = create_response.json()["id"]
+
+    # Now, delete the thread
+    delete_response = requests.delete(
+        f"{THREADS_URL}/{thread_id}",
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert delete_response.status_code == 200
+
+    response_json = delete_response.json()
+    assert response_json["id"] == thread_id
+    assert response_json["object"] == "thread.deleted"
+    assert response_json["deleted"] is True
+
+
+def test_list_threads(admin_user: DATestUser | None) -> None:
+    # Create a few threads
+    for _ in range(3):
+        requests.post(
+            THREADS_URL,
+            json={},
+            headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+        )
+
+    # Now, list the threads
+    list_response = requests.get(
+        THREADS_URL,
+        headers=admin_user.headers if admin_user else GENERAL_HEADERS,
+    )
+    assert list_response.status_code == 200
+
+    response_json = list_response.json()
+    assert "sessions" in response_json
+    assert len(response_json["sessions"]) >= 3
+
+    for session in response_json["sessions"]:
+        assert "id" in session
+        assert "name" in session
+        assert "persona_id" in session
+        assert "time_created" in session
+        assert "shared_status" in session
+        assert "folder_id" in session
+        assert "current_alternate_model" in session
+
+        # Validate UUID
+        UUID(session["id"])
+
+        # Validate shared_status
+        assert session["shared_status"] in [
+            status.value for status in ChatSessionSharedStatus
+        ]

backend/tests/integration/tests/connector/test_connector_deletion.py

@@ -29,6 +29,78 @@ from tests.integration.common_utils.test_models import DATestUserGroup
 from tests.integration.common_utils.vespa import vespa_fixture
 
 
+# def test_connector_creation(reset: None) -> None:
+#     # Creating an admin user (first user created is automatically an admin)
+#     admin_user: DATestUser = UserManager.create(name="admin_user")
+
+#     # create connectors
+#     cc_pair_1 = CCPairManager.create_from_scratch(
+#         source=DocumentSource.INGESTION_API,
+#         user_performing_action=admin_user,
+#     )
+
+#     cc_pair_info = CCPairManager.get_single(
+#         cc_pair_1.id, user_performing_action=admin_user
+#     )
+#     assert cc_pair_info
+#     assert cc_pair_info.creator
+#     assert str(cc_pair_info.creator) == admin_user.id
+#     assert cc_pair_info.creator_email == admin_user.email
+
+
+# TODO(rkuo): will enable this once i have credentials on github
+# def test_overlapping_connector_creation(reset: None) -> None:
+#     # Creating an admin user (first user created is automatically an admin)
+#     admin_user: DATestUser = UserManager.create(name="admin_user")
+
+#     config = {
+#         "wiki_base": os.environ["CONFLUENCE_TEST_SPACE_URL"],
+#         "space": os.environ["CONFLUENCE_TEST_SPACE"],
+#         "is_cloud": True,
+#         "page_id": "",
+#     }
+
+#     credential = {
+#         "confluence_username": os.environ["CONFLUENCE_USER_NAME"],
+#         "confluence_access_token": os.environ["CONFLUENCE_ACCESS_TOKEN"],
+#     }
+
+#     # store the time before we create the connector so that we know after
+#     # when the indexing should have started
+#     now = datetime.now(timezone.utc)
+
+#     # create connector
+#     cc_pair_1 = CCPairManager.create_from_scratch(
+#         source=DocumentSource.CONFLUENCE,
+#         connector_specific_config=config,
+#         credential_json=credential,
+#         user_performing_action=admin_user,
+#     )
+
+#     CCPairManager.wait_for_indexing(
+#         cc_pair_1, now, timeout=60, user_performing_action=admin_user
+#     )
+
+#     cc_pair_2 = CCPairManager.create_from_scratch(
+#         source=DocumentSource.CONFLUENCE,
+#         connector_specific_config=config,
+#         credential_json=credential,
+#         user_performing_action=admin_user,
+#     )
+
+#     CCPairManager.wait_for_indexing(
+#         cc_pair_2, now, timeout=60, user_performing_action=admin_user
+#     )
+
+#     info_1 = CCPairManager.get_single(cc_pair_1.id)
+#     assert info_1
+
+#     info_2 = CCPairManager.get_single(cc_pair_2.id)
+#     assert info_2
+
+#     assert info_1.num_docs_indexed == info_2.num_docs_indexed
+
+
 def test_connector_deletion(reset: None, vespa_client: vespa_fixture) -> None:
     # Creating an admin user (first user created is automatically an admin)
     admin_user: DATestUser = UserManager.create(name="admin_user")

ct.yaml

@@ -1,12 +1,18 @@
 # See https://github.com/helm/chart-testing#configuration
 
+# still have to specify this on the command line for list-changed
 chart-dirs:
   - deployment/helm/charts
 
+# must be kept in sync with Chart.yaml
 chart-repos:
-  - vespa=https://unoplat.github.io/vespa-helm-charts
+  - vespa=https://danswer-ai.github.io/vespa-helm-charts
   - postgresql=https://charts.bitnami.com/bitnami
   
-helm-extra-args: --timeout 600s
+helm-extra-args: --debug --timeout 600s
+
+# nginx appears to not work on kind, likely due to lack of loadbalancer support
+# helm-extra-set-args also only works on the command line, not in this yaml
+# helm-extra-set-args: --set=nginx.enabled=false
 
 validate-maintainers: false

deployment/helm/charts/danswer/Chart.lock

@@ -3,13 +3,13 @@ dependencies:
   repository: https://charts.bitnami.com/bitnami
   version: 14.3.1
 - name: vespa
-  repository: https://unoplat.github.io/vespa-helm-charts
-  version: 0.2.3
+  repository: https://danswer-ai.github.io/vespa-helm-charts
+  version: 0.2.16
 - name: nginx
   repository: oci://registry-1.docker.io/bitnamicharts
   version: 15.14.0
 - name: redis
   repository: https://charts.bitnami.com/bitnami
   version: 20.1.0
-digest: sha256:fb42426c1d13667a4929d0d6a7d681bf08120e4a4eb1d15437e4ec70920be3f8
-generated: "2024-09-11T09:16:03.312328-07:00"
+digest: sha256:711bbb76ba6ab604a36c9bf1839ab6faa5610afb21e535afd933c78f2d102232
+generated: "2024-11-07T09:39:30.17171-08:00"

deployment/helm/charts/danswer/Chart.yaml

@@ -5,7 +5,7 @@ home: https://www.danswer.ai/
 sources:
   - "https://github.com/danswer-ai/danswer"
 type: application
-version: 0.2.0
+version: 0.2.1
 appVersion: "latest"
 annotations:
   category: Productivity
@@ -23,8 +23,8 @@ dependencies:
     repository: https://charts.bitnami.com/bitnami
     condition: postgresql.enabled
   - name: vespa
-    version: 0.2.3
-    repository: https://unoplat.github.io/vespa-helm-charts
+    version: 0.2.16
+    repository: https://danswer-ai.github.io/vespa-helm-charts
     condition: vespa.enabled
   - name: nginx
     version: 15.14.0

deployment/helm/charts/danswer/templates/configmap.yaml

@@ -7,7 +7,7 @@ metadata:
 data:
   INTERNAL_URL: "http://{{ include "danswer-stack.fullname" . }}-api-service:{{ .Values.api.service.port | default 8080 }}"
   POSTGRES_HOST: {{ .Release.Name }}-postgresql
-  VESPA_HOST: "document-index-service"
+  VESPA_HOST: da-vespa-0.vespa-service
   REDIS_HOST: {{ .Release.Name }}-redis-master
   MODEL_SERVER_HOST: "{{ include "danswer-stack.fullname" . }}-inference-model-service"
   INDEXING_MODEL_SERVER_HOST: "{{ include "danswer-stack.fullname" . }}-indexing-model-service"

deployment/helm/charts/danswer/templates/tests/test-connection.yaml

@@ -11,5 +11,5 @@ spec:
     - name: wget
       image: busybox
       command: ['wget']
-      args: ['{{ include "danswer-stack.fullname" . }}:{{ .Values.webserver.service.port }}']
+      args: ['{{ include "danswer-stack.fullname" . }}-webserver:{{ .Values.webserver.service.port }}']
   restartPolicy: Never

examples/assistants-api/topics_analyzer.py

@@ -0,0 +1,125 @@
+import argparse
+import os
+import time
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
+
+from openai import OpenAI
+
+
+ASSISTANT_NAME = "Topic Analyzer"
+SYSTEM_PROMPT = """
+You are a helpful assistant that analyzes topics by searching through available \
+documents and providing insights. These available documents come from common \
+workplace tools like Slack, emails, Confluence, Google Drive, etc.
+
+When analyzing a topic:
+1. Search for relevant information using the search tool
+2. Synthesize the findings into clear insights
+3. Highlight key trends, patterns, or notable developments
+4. Maintain objectivity and cite sources where relevant
+"""
+USER_PROMPT = """
+Please analyze and provide insights about this topic: {topic}.
+
+IMPORTANT: do not mention things that are not relevant to the specified topic. \
+If there is no relevant information, just say "No relevant information found."
+"""
+
+
+def wait_on_run(client: OpenAI, run, thread):
+    while run.status == "queued" or run.status == "in_progress":
+        run = client.beta.threads.runs.retrieve(
+            thread_id=thread.id,
+            run_id=run.id,
+        )
+        time.sleep(0.5)
+    return run
+
+
+def show_response(messages) -> None:
+    # Get only the assistant's response text
+    for message in messages.data[::-1]:
+        if message.role == "assistant":
+            for content in message.content:
+                if content.type == "text":
+                    print(content.text)
+                    break
+
+
+def analyze_topics(topics: list[str]) -> None:
+    openai_api_key = os.environ.get(
+        "OPENAI_API_KEY", "<your OpenAI API key if not set as env var>"
+    )
+    danswer_api_key = os.environ.get(
+        "DANSWER_API_KEY", "<your Danswer API key if not set as env var>"
+    )
+    client = OpenAI(
+        api_key=openai_api_key,
+        base_url="http://localhost:8080/openai-assistants",
+        default_headers={
+            "Authorization": f"Bearer {danswer_api_key}",
+        },
+    )
+
+    # Create an assistant if it doesn't exist
+    try:
+        assistants = client.beta.assistants.list(limit=100)
+        # Find the Topic Analyzer assistant if it exists
+        assistant = next((a for a in assistants.data if a.name == ASSISTANT_NAME))
+        client.beta.assistants.delete(assistant.id)
+    except Exception:
+        pass
+
+    assistant = client.beta.assistants.create(
+        name=ASSISTANT_NAME,
+        instructions=SYSTEM_PROMPT,
+        tools=[{"type": "SearchTool"}],  # type: ignore
+        model="gpt-4o",
+    )
+
+    # Process each topic individually
+    for topic in topics:
+        thread = client.beta.threads.create()
+        message = client.beta.threads.messages.create(
+            thread_id=thread.id,
+            role="user",
+            content=USER_PROMPT.format(topic=topic),
+        )
+
+        run = client.beta.threads.runs.create(
+            thread_id=thread.id,
+            assistant_id=assistant.id,
+            tools=[
+                {  # type: ignore
+                    "type": "SearchTool",
+                    "retrieval_details": {
+                        "run_search": "always",
+                        "filters": {
+                            "time_cutoff": str(
+                                datetime.now(timezone.utc) - timedelta(days=7)
+                            )
+                        },
+                    },
+                }
+            ],
+        )
+
+        run = wait_on_run(client, run, thread)
+        messages = client.beta.threads.messages.list(
+            thread_id=thread.id, order="asc", after=message.id
+        )
+        print(f"\nAnalysis for topic: {topic}")
+        print("-" * 40)
+        show_response(messages)
+        print()
+
+
+# Example usage
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Analyze specific topics")
+    parser.add_argument("topics", nargs="+", help="Topics to analyze (one or more)")
+
+    args = parser.parse_args()
+    analyze_topics(args.topics)

web/src/app/admin/indexing/status/CCPairIndexingStatusTable.tsx

@@ -194,7 +194,9 @@ function ConnectorRow({
   return (
     <TableRow
       className={`hover:bg-hover-light ${
-        invisible ? "invisible !h-0 !-mb-10" : "!border !border-border"
+        invisible
+          ? "invisible !h-0 !-mb-10 !border-none"
+          : "!border !border-border"
       }  w-full cursor-pointer relative `}
       onClick={() => {
         router.push(`/admin/connector/${ccPairsIndexingStatus.cc_pair_id}`);
@@ -434,7 +436,6 @@ export function CCPairIndexingStatusTable({
           {!shouldExpand ? "Collapse All" : "Expand All"}
         </Button>
       </div>
-
       <TableBody>
         {sortedSources
           .filter(
@@ -454,14 +455,12 @@ export function CCPairIndexingStatusTable({
               return (
                 <React.Fragment key={ind}>
                   <br className="mt-4" />
-
                   <SummaryRow
                     source={source}
                     summary={groupSummaries[source]}
                     isOpen={connectorsToggled[source] || false}
                     onToggle={() => toggleSource(source)}
                   />
-
                   {connectorsToggled[source] && (
                     <>
                       <TableRow className="border border-border">

web/src/app/admin/users/page.tsx

@@ -188,7 +188,7 @@ const AddUserButton = ({
   };
   return (
     <>
-      <Button className="w-fit" onClick={() => setModal(true)}>
+      <Button className="my-auto w-fit" onClick={() => setModal(true)}>
         <div className="flex">
           <FiPlusSquare className="my-auto mr-2" />
           Invite Users

web/src/app/assistants/gallery/page.tsx

@@ -4,7 +4,6 @@ import { fetchChatData } from "@/lib/chat/fetchChatData";
 import { unstable_noStore as noStore } from "next/cache";
 import { redirect } from "next/navigation";
 import WrappedAssistantsGallery from "./WrappedAssistantsGallery";
-import { AssistantsProvider } from "@/components/context/AssistantsContext";
 import { cookies } from "next/headers";
 
 export default async function GalleryPage(props: {

web/src/app/chat/ChatPage.tsx

@@ -142,6 +142,16 @@ export function ChatPage({
     refreshChatSessions,
   } = useChatContext();
 
+  // handle redirect if chat page is disabled
+  // NOTE: this must be done here, in a client component since
+  // settings are passed in via Context and therefore aren't
+  // available in server-side components
+  const settings = useContext(SettingsContext);
+  const enterpriseSettings = settings?.enterpriseSettings;
+  if (settings?.settings?.chat_page_enabled === false) {
+    router.push("/search");
+  }
+
   const { assistants: availableAssistants, finalAssistants } = useAssistants();
 
   const [showApiKeyModal, setShowApiKeyModal] = useState(
@@ -881,7 +891,6 @@ export function ChatPage({
     }, 1500);
   };
 
-  const distance = 500; // distance that should "engage" the scroll
   const debounceNumber = 100; // time for debouncing
 
   const [hasPerformedInitialScroll, setHasPerformedInitialScroll] = useState(
@@ -1545,17 +1554,6 @@ export function ChatPage({
       }
     });
   };
-
-  // handle redirect if chat page is disabled
-  // NOTE: this must be done here, in a client component since
-  // settings are passed in via Context and therefore aren't
-  // available in server-side components
-  const settings = useContext(SettingsContext);
-  const enterpriseSettings = settings?.enterpriseSettings;
-  if (settings?.settings?.chat_page_enabled === false) {
-    router.push("/search");
-  }
-
   const [showDocSidebar, setShowDocSidebar] = useState(false); // State to track if sidebar is open
 
   // Used to maintain a "time out" for history sidebar so our existing refs can have time to process change
@@ -1603,9 +1601,9 @@ export function ChatPage({
     scrollableDivRef,
     scrollDist,
     endDivRef,
-    distance,
     debounceNumber,
     waitForScrollRef,
+    mobile: settings?.isMobile,
   });
 
   // Virtualization + Scrolling related effects and functions

web/src/app/chat/lib.tsx

@@ -639,19 +639,22 @@ export async function useScrollonStream({
   scrollableDivRef,
   scrollDist,
   endDivRef,
-  distance,
   debounceNumber,
-  waitForScrollRef,
+  mobile,
 }: {
   chatState: ChatState;
   scrollableDivRef: RefObject<HTMLDivElement>;
   waitForScrollRef: RefObject<boolean>;
   scrollDist: MutableRefObject<number>;
   endDivRef: RefObject<HTMLDivElement>;
-  distance: number;
   debounceNumber: number;
   mobile?: boolean;
 }) {
+  const mobileDistance = 900; // distance that should "engage" the scroll
+  const desktopDistance = 500; // distance that should "engage" the scroll
+
+  const distance = mobile ? mobileDistance : desktopDistance;
+
   const preventScrollInterference = useRef<boolean>(false);
   const preventScroll = useRef<boolean>(false);
   const blockActionRef = useRef<boolean>(false);
@@ -692,7 +695,7 @@ export async function useScrollonStream({
         endDivRef.current
       ) {
         // catch up if necessary!
-        const scrollAmount = scrollDist.current + 10000;
+        const scrollAmount = scrollDist.current + (mobile ? 1000 : 10000);
         if (scrollDist.current > 300) {
           // if (scrollDist.current > 140) {
           endDivRef.current.scrollIntoView();

web/src/components/admin/ClientLayout.tsx

@@ -420,7 +420,7 @@ export function ClientLayout({
           <div className="fixed bg-background left-0 gap-x-4 mb-8 px-4 py-2 w-full items-center flex justify-end">
             <UserDropdown />
           </div>
-          <div className="pt-20 flex overflow-y-auto h-full px-4 md:px-12">
+          <div className="pt-20 flex overflow-y-auto overflow-x-hidden h-full px-4 md:px-12">
             {children}
           </div>
         </div>

web/src/components/ui/button.tsx

@@ -12,7 +12,7 @@ const buttonVariants = cva(
         success:
           "bg-green-100 text-green-600 hover:bg-green-500/90 dark:bg-blue-500 dark:text-neutral-50 dark:hover:bg-green-900/90",
         "success-reverse":
-          "bg-green-500 text-white hover:bg-green-600/90 dark:bg-neutral-50 dark:text-blue-500 dark:hover:bg-green-100/90",
+          "bg-green-500 text-inverted hover:bg-green-600/90 dark:bg-neutral-50 dark:text-blue-500 dark:hover:bg-green-100/90",
 
         default:
           "bg-neutral-900 border-border text-neutral-50 hover:bg-neutral-900/90 dark:bg-neutral-50 dark:text-neutral-900 dark:hover:bg-neutral-50/90",
@@ -38,7 +38,7 @@ const buttonVariants = cva(
         "link-reverse":
           "text-neutral-50 underline-offset-4 hover:underline dark:text-neutral-900",
         submit:
-          "bg-green-500 text-green-100 hover:bg-green-600/90 dark:bg-neutral-50 dark:text-blue-500 dark:hover:bg-green-100/90",
+          "bg-green-500 text-inverted hover:bg-green-600/90 dark:bg-neutral-50 dark:text-blue-500 dark:hover:bg-green-100/90",
 
         // "bg-blue-600 text-neutral-50 hover:bg-blue-600/80 dark:bg-blue-600 dark:text-neutral-50 dark:hover:bg-blue-600/90",
         "submit-reverse":