mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-22 11:22:49 -05:00
Chris Weaver
0c58c8d6cb
Adds: - name for connector credential pairs + frontend changes to start populating this field - document set table migration - during indexing, document sets are now checked and inserted into Vespa - background job to check if document sets need to be synced - document set management APIs - document set management dashboard in the UI
54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
from celery.result import AsyncResult
|
|
from sqlalchemy.orm import Session
|
|
|
|
from danswer.background.celery.celery import sync_document_set_task
|
|
from danswer.background.utils import interval_run_job
|
|
from danswer.db.document_set import (
|
|
fetch_document_sets,
|
|
)
|
|
from danswer.db.engine import get_sqlalchemy_engine
|
|
from danswer.utils.logger import setup_logger
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
_ExistingTaskCache: dict[int, AsyncResult] = {}
|
|
|
|
|
|
def _document_sync_loop() -> None:
|
|
# cleanup tasks
|
|
existing_tasks = list(_ExistingTaskCache.items())
|
|
for document_set_id, task in existing_tasks:
|
|
if task.ready():
|
|
logger.info(
|
|
f"Document set '{document_set_id}' is complete with status "
|
|
f"{task.status}. Cleaning up."
|
|
)
|
|
del _ExistingTaskCache[document_set_id]
|
|
|
|
# kick off new tasks
|
|
with Session(get_sqlalchemy_engine()) as db_session:
|
|
# check if any document sets are not synced
|
|
document_set_info = fetch_document_sets(db_session=db_session)
|
|
for document_set, _ in document_set_info:
|
|
if not document_set.is_up_to_date:
|
|
if document_set.id in _ExistingTaskCache:
|
|
logger.info(
|
|
f"Document set '{document_set.id}' is already syncing. Skipping."
|
|
)
|
|
continue
|
|
|
|
logger.info(
|
|
f"Document set {document_set.id} is not synced. Syncing now!"
|
|
)
|
|
task = sync_document_set_task.apply_async(
|
|
kwargs=dict(document_set_id=document_set.id),
|
|
)
|
|
_ExistingTaskCache[document_set.id] = task
|
|
|
|
|
|
if __name__ == "__main__":
|
|
interval_run_job(
|
|
job=_document_sync_loop, delay=5, emit_job_start_log=False
|
|
) # run every 5 seconds
|