tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-06 19:25:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

refact: add batch size for scan cmd

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
memos/cmds/library.py
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 updated_file_count = 0
added_file_count = 0 added_file_count = 0
scanned_files = set() scanned_files = set()
semaphore = asyncio.Semaphore(settings.batchsize) semaphore = asyncio.Semaphore(batch_size)
async with httpx.AsyncClient(timeout=60) as client: async with httpx.AsyncClient(timeout=60) as client:
tasks = [] tasks = []
@ -377,6 +377,7 @@ def scan(
force: bool = False, force: bool = False,
plugins: List[int] = typer.Option(None, "--plugin", "-p"), plugins: List[int] = typer.Option(None, "--plugin", "-p"),
folders: List[int] = typer.Option(None, "--folder", "-f"), 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 # Check if both path and folders are provided
if path and folders: if path and folders:
@ -426,7 +427,7 @@ def scan(
continue continue
added_file_count, updated_file_count, scanned_files = asyncio.run( 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_added += added_file_count
total_files_updated += updated_file_count total_files_updated += updated_file_count

11
memos/commands.py
View File

@ -106,6 +106,7 @@ def get_or_create_default_library():
Ensure the library has at least one folder. Ensure the library has at least one folder.
""" """
from .cmds.plugin import bind from .cmds.plugin import bind
response = httpx.get(f"{BASE_URL}/libraries") response = httpx.get(f"{BASE_URL}/libraries")
if response.status_code != 200: if response.status_code != 200:
print(f"Failed to retrieve libraries: {response.status_code} - {response.text}") 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"), path: str = typer.Argument(None, help="Path to scan within the library"),
plugins: List[int] = typer.Option(None, "--plugin", "-p"), plugins: List[int] = typer.Option(None, "--plugin", "-p"),
folders: List[int] = typer.Option(None, "--folder", "-f"), 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. 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']}") print(f"Scanning library: {default_library['name']}")
scan( 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,
) )

2
memos/config.py
View File

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

1
memos/default_config.yaml
View File

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