mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-22 11:22:49 -05:00
85 lines
3.3 KiB
Python
85 lines
3.3 KiB
Python
from typing import cast
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from danswer.datastores.document_index import get_default_document_index
|
|
from danswer.datastores.interfaces import DocumentIndex
|
|
from danswer.datastores.interfaces import UpdateRequest
|
|
from danswer.db.document import prepare_to_modify_documents
|
|
from danswer.db.document_set import delete_document_set
|
|
from danswer.db.document_set import fetch_document_sets_for_documents
|
|
from danswer.db.document_set import fetch_documents_for_document_set
|
|
from danswer.db.document_set import get_document_set_by_id
|
|
from danswer.db.document_set import mark_document_set_as_synced
|
|
from danswer.db.engine import get_sqlalchemy_engine
|
|
from danswer.db.models import DocumentSet
|
|
from danswer.utils.batching import batch_generator
|
|
from danswer.utils.logger import setup_logger
|
|
|
|
logger = setup_logger()
|
|
|
|
_SYNC_BATCH_SIZE = 1000
|
|
|
|
|
|
def _sync_document_batch(
|
|
document_ids: list[str], document_index: DocumentIndex
|
|
) -> None:
|
|
logger.debug(f"Syncing document sets for: {document_ids}")
|
|
# begin a transaction, release lock at the end
|
|
with Session(get_sqlalchemy_engine()) as db_session:
|
|
# acquires a lock on the documents so that no other process can modify them
|
|
prepare_to_modify_documents(db_session=db_session, document_ids=document_ids)
|
|
|
|
# get current state of document sets for these documents
|
|
document_set_map = {
|
|
document_id: document_sets
|
|
for document_id, document_sets in fetch_document_sets_for_documents(
|
|
document_ids=document_ids, db_session=db_session
|
|
)
|
|
}
|
|
|
|
# update Vespa
|
|
document_index.update(
|
|
update_requests=[
|
|
UpdateRequest(
|
|
document_ids=[document_id],
|
|
document_sets=set(document_set_map.get(document_id, [])),
|
|
)
|
|
for document_id in document_ids
|
|
]
|
|
)
|
|
|
|
|
|
def sync_document_set(document_set_id: int) -> None:
|
|
document_index = get_default_document_index()
|
|
with Session(get_sqlalchemy_engine()) as db_session:
|
|
documents_to_update = fetch_documents_for_document_set(
|
|
document_set_id=document_set_id,
|
|
db_session=db_session,
|
|
current_only=False,
|
|
)
|
|
for document_batch in batch_generator(documents_to_update, _SYNC_BATCH_SIZE):
|
|
_sync_document_batch(
|
|
document_ids=[document.id for document in document_batch],
|
|
document_index=document_index,
|
|
)
|
|
|
|
# if there are no connectors, then delete the document set. Otherwise, just
|
|
# mark it as successfully synced.
|
|
document_set = cast(
|
|
DocumentSet,
|
|
get_document_set_by_id(
|
|
db_session=db_session, document_set_id=document_set_id
|
|
),
|
|
) # casting since we "know" a document set with this ID exists
|
|
if not document_set.connector_credential_pairs:
|
|
delete_document_set(document_set_row=document_set, db_session=db_session)
|
|
logger.info(
|
|
f"Successfully deleted document set with ID: '{document_set_id}'!"
|
|
)
|
|
else:
|
|
mark_document_set_as_synced(
|
|
document_set_id=document_set_id, db_session=db_session
|
|
)
|
|
logger.info(f"Document set sync for '{document_set_id}' complete!")
|