fix: use len(rows) < batch_size early-exit instead of total-count loop bound
- Replace 'while offset < count/total' with 'while True' + break on short batch - Fixes tool_list_rooms iterating over unfiltered col.count() when wing filter active - Fixes all 4 paginated functions (tool_status, tool_list_wings, tool_list_rooms, tool_get_taxonomy) missing early-exit when batch smaller than batch_size - Remove unused 'total' variable in tool_list_wings, tool_list_rooms, tool_get_taxonomy (replaced col.count() with accessibility check only) Per bensig review comments on PR #371
This commit is contained in:
+27
-15
@@ -146,15 +146,18 @@ def tool_status():
|
||||
batch_size = 5000
|
||||
offset = 0
|
||||
error_info = None
|
||||
while offset < count:
|
||||
while True:
|
||||
try:
|
||||
batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
|
||||
for m in batch["metadatas"]:
|
||||
rows = batch["metadatas"]
|
||||
for m in rows:
|
||||
w = m.get("wing", "unknown")
|
||||
r = m.get("room", "unknown")
|
||||
wings[w] = wings.get(w, 0) + 1
|
||||
rooms[r] = rooms.get(r, 0) + 1
|
||||
offset += batch_size
|
||||
offset += len(rows)
|
||||
if len(rows) < batch_size:
|
||||
break
|
||||
except Exception as e:
|
||||
error_info = f"Partial result, failed at offset {offset}: {str(e)}"
|
||||
break
|
||||
@@ -213,16 +216,19 @@ def tool_list_wings():
|
||||
batch_size = 5000
|
||||
offset = 0
|
||||
try:
|
||||
total = col.count()
|
||||
col.count() # verify collection is accessible
|
||||
except Exception as e:
|
||||
return {"wings": {}, "error": str(e)}
|
||||
while offset < total:
|
||||
while True:
|
||||
try:
|
||||
batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
|
||||
for m in batch["metadatas"]:
|
||||
rows = batch["metadatas"]
|
||||
for m in rows:
|
||||
w = m.get("wing", "unknown")
|
||||
wings[w] = wings.get(w, 0) + 1
|
||||
offset += batch_size
|
||||
offset += len(rows)
|
||||
if len(rows) < batch_size:
|
||||
break
|
||||
except Exception as e:
|
||||
return {
|
||||
"wings": wings,
|
||||
@@ -241,19 +247,22 @@ def tool_list_rooms(wing: str = None):
|
||||
offset = 0
|
||||
where = {"wing": wing} if wing else None
|
||||
try:
|
||||
total = col.count()
|
||||
col.count() # verify collection is accessible
|
||||
except Exception as e:
|
||||
return {"wing": wing or "all", "rooms": {}, "error": str(e)}
|
||||
while offset < total:
|
||||
while True:
|
||||
try:
|
||||
kwargs = {"include": ["metadatas"], "limit": batch_size, "offset": offset}
|
||||
if where:
|
||||
kwargs["where"] = where
|
||||
batch = col.get(**kwargs)
|
||||
for m in batch["metadatas"]:
|
||||
rows = batch["metadatas"]
|
||||
for m in rows:
|
||||
r = m.get("room", "unknown")
|
||||
rooms[r] = rooms.get(r, 0) + 1
|
||||
offset += batch_size
|
||||
offset += len(rows)
|
||||
if len(rows) < batch_size:
|
||||
break
|
||||
except Exception as e:
|
||||
return {
|
||||
"wing": wing or "all",
|
||||
@@ -272,19 +281,22 @@ def tool_get_taxonomy():
|
||||
batch_size = 5000
|
||||
offset = 0
|
||||
try:
|
||||
total = col.count()
|
||||
col.count() # verify collection is accessible
|
||||
except Exception as e:
|
||||
return {"taxonomy": {}, "error": str(e)}
|
||||
while offset < total:
|
||||
while True:
|
||||
try:
|
||||
batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
|
||||
for m in batch["metadatas"]:
|
||||
rows = batch["metadatas"]
|
||||
for m in rows:
|
||||
w = m.get("wing", "unknown")
|
||||
r = m.get("room", "unknown")
|
||||
if w not in taxonomy:
|
||||
taxonomy[w] = {}
|
||||
taxonomy[w][r] = taxonomy[w].get(r, 0) + 1
|
||||
offset += batch_size
|
||||
offset += len(rows)
|
||||
if len(rows) < batch_size:
|
||||
break
|
||||
except Exception as e:
|
||||
return {
|
||||
"taxonomy": taxonomy,
|
||||
|
||||
Reference in New Issue
Block a user