refact: add batch size for scan cmd

This commit is contained in:
arkohut 2024-11-06 16:41:10 +08:00
parent b387807ee4
commit f524622728
4 changed files with 18 additions and 11 deletions

View File

@ -169,11 +169,11 @@ def is_temp_file(filename):
)
async def loop_files(library_id, folder, folder_path, force, plugins):
async def loop_files(library_id, folder, folder_path, force, plugins, batch_size):
updated_file_count = 0
added_file_count = 0
scanned_files = set()
semaphore = asyncio.Semaphore(settings.batchsize)
semaphore = asyncio.Semaphore(batch_size)
async with httpx.AsyncClient(timeout=60) as client:
tasks = []
@ -377,6 +377,7 @@ def scan(
force: bool = False,
plugins: List[int] = typer.Option(None, "--plugin", "-p"),
folders: List[int] = typer.Option(None, "--folder", "-f"),
batch_size: int = typer.Option(1, "--batch-size", "-bs", help="Batch size for processing files"),
):
# Check if both path and folders are provided
if path and folders:
@ -426,7 +427,7 @@ def scan(
continue
added_file_count, updated_file_count, scanned_files = asyncio.run(
loop_files(library_id, folder, folder_path, force, plugins)
loop_files(library_id, folder, folder_path, force, plugins, batch_size)
)
total_files_added += added_file_count
total_files_updated += updated_file_count

View File

@ -106,6 +106,7 @@ def get_or_create_default_library():
Ensure the library has at least one folder.
"""
from .cmds.plugin import bind
response = httpx.get(f"{BASE_URL}/libraries")
if response.status_code != 200:
print(f"Failed to retrieve libraries: {response.status_code} - {response.text}")
@ -162,6 +163,9 @@ def scan_default_library(
path: str = typer.Argument(None, help="Path to scan within the library"),
plugins: List[int] = typer.Option(None, "--plugin", "-p"),
folders: List[int] = typer.Option(None, "--folder", "-f"),
batch_size: int = typer.Option(
1, "--batch-size", "-bs", help="Batch size for processing files"
),
):
"""
Scan the screenshots directory and add it to the library if empty.
@ -174,7 +178,12 @@ def scan_default_library(
print(f"Scanning library: {default_library['name']}")
scan(
default_library["id"], path=path, plugins=plugins, folders=folders, force=force
default_library["id"],
path=path,
plugins=plugins,
folders=folders,
force=force,
batch_size=batch_size,
)

View File

@ -71,8 +71,6 @@ class Settings(BaseSettings):
# Embedding settings
embedding: EmbeddingSettings = EmbeddingSettings()
batchsize: int = 1
auth_username: str = "admin"
auth_password: SecretStr = SecretStr("changeme")

View File

@ -6,7 +6,6 @@ screenshots_dir: screenshots
server_host: 0.0.0.0
server_port: 8839
batchsize: 1
auth_username: admin
auth_password: changeme
default_plugins: