enable batch

This commit is contained in:
Billy Bao 2025-12-11 19:55:39 +08:00
parent 97de53586c
commit fa54b955e6
2 changed files with 2 additions and 3 deletions

View file

@ -14,10 +14,11 @@ from common.data_source.models import Document, GenerateDocumentsOutput
from common.data_source.utils import get_file_ext
class BoxConnector(LoadConnector, PollConnector):
def __init__(self, folder_id: str, batch_size: int = INDEX_BATCH_SIZE, use_marker: bool = False) -> None:
def __init__(self, folder_id: str, batch_size: int = INDEX_BATCH_SIZE, use_marker: bool = True) -> None:
self.batch_size = batch_size
self.folder_id = "0" if not folder_id else folder_id
self.use_marker = use_marker
def load_credentials(self, auth: Any):
self.box_client = BoxClient(auth=auth)

View file

@ -634,9 +634,7 @@ class BOX(SyncBase):
async def _generate(self, task: dict):
self.connector = BoxConnector(
batch_size=self.conf.get("batch_size", INDEX_BATCH_SIZE),
folder_id=self.conf.get("folder_id", "0"),
use_marker=self.conf.get("use_marker", False)
)
credential = json.loads(self.conf['credentials']['box_tokens'])