tcsenpai/pensieve
1
0
Fork 0
mirror of https://github.com/tcsenpai/pensieve.git synced 2025-06-10 13:07:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: show full progress when index a folder

Browse Source
This commit is contained in:
arkohut 2024-07-29 16:50:40 +08:00
parent 367b818992
commit 2794e6b7b4
1 changed files with 39 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
memos/commands.py
View File

@ -447,15 +447,20 @@ def index(
tqdm.write(f"Processing folder: {folder['id']}")
# List all entities in the folder
limit = 200
offset = 0
total_entities = 0 # We'll update this after the first request
with tqdm(
total=total_entities, desc="Indexing entities", leave=True
) as pbar:
while True:
entities_response = httpx.get(
f"{BASE_URL}/libraries/{library_id}/folders/{folder['id']}/entities",
params={"limit": 200, "offset": offset},
params={"limit": limit, "offset": offset},
timeout=60,
)
if entities_response.status_code != 200:
tqdm.write(
pbar.write(
f"Failed to get entities: {entities_response.status_code} - {entities_response.text}"
)
break
@ -464,19 +469,28 @@ def index(
if not entities:
break
# Update total if this is the first request
if offset == 0:
total_entities = int(
entities_response.headers.get("X-Total-Count", total_entities)
)
pbar.total = total_entities
pbar.refresh()
# Index each entity
for entity in tqdm(entities, desc="Indexing entities", leave=False):
for entity in entities:
index_response = httpx.post(f"{BASE_URL}/entities/{entity['id']}/index")
if index_response.status_code == 204:
tqdm.write(f"Indexed entity: {entity['id']}")
pbar.write(f"Indexed entity: {entity['id']}")
else:
tqdm.write(
pbar.write(
f"Failed to index entity {entity['id']}: {index_response.status_code} - {index_response.text}"
)
scanned_entities.add(str(entity["id"]))
pbar.update(1)
offset += 200
offset += limit
# List all indexed entities in the folder
offset = 0