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:
RhettOP
2026-04-10 17:15:36 +01:00
parent a3fec4f565
commit 8a6e75eed8
+27 -15
View File
@@ -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,