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
|
batch_size = 5000
|
||||||
offset = 0
|
offset = 0
|
||||||
error_info = None
|
error_info = None
|
||||||
while offset < count:
|
while True:
|
||||||
try:
|
try:
|
||||||
batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
|
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")
|
w = m.get("wing", "unknown")
|
||||||
r = m.get("room", "unknown")
|
r = m.get("room", "unknown")
|
||||||
wings[w] = wings.get(w, 0) + 1
|
wings[w] = wings.get(w, 0) + 1
|
||||||
rooms[r] = rooms.get(r, 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:
|
except Exception as e:
|
||||||
error_info = f"Partial result, failed at offset {offset}: {str(e)}"
|
error_info = f"Partial result, failed at offset {offset}: {str(e)}"
|
||||||
break
|
break
|
||||||
@@ -213,16 +216,19 @@ def tool_list_wings():
|
|||||||
batch_size = 5000
|
batch_size = 5000
|
||||||
offset = 0
|
offset = 0
|
||||||
try:
|
try:
|
||||||
total = col.count()
|
col.count() # verify collection is accessible
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return {"wings": {}, "error": str(e)}
|
return {"wings": {}, "error": str(e)}
|
||||||
while offset < total:
|
while True:
|
||||||
try:
|
try:
|
||||||
batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
|
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")
|
w = m.get("wing", "unknown")
|
||||||
wings[w] = wings.get(w, 0) + 1
|
wings[w] = wings.get(w, 0) + 1
|
||||||
offset += batch_size
|
offset += len(rows)
|
||||||
|
if len(rows) < batch_size:
|
||||||
|
break
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return {
|
return {
|
||||||
"wings": wings,
|
"wings": wings,
|
||||||
@@ -241,19 +247,22 @@ def tool_list_rooms(wing: str = None):
|
|||||||
offset = 0
|
offset = 0
|
||||||
where = {"wing": wing} if wing else None
|
where = {"wing": wing} if wing else None
|
||||||
try:
|
try:
|
||||||
total = col.count()
|
col.count() # verify collection is accessible
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return {"wing": wing or "all", "rooms": {}, "error": str(e)}
|
return {"wing": wing or "all", "rooms": {}, "error": str(e)}
|
||||||
while offset < total:
|
while True:
|
||||||
try:
|
try:
|
||||||
kwargs = {"include": ["metadatas"], "limit": batch_size, "offset": offset}
|
kwargs = {"include": ["metadatas"], "limit": batch_size, "offset": offset}
|
||||||
if where:
|
if where:
|
||||||
kwargs["where"] = where
|
kwargs["where"] = where
|
||||||
batch = col.get(**kwargs)
|
batch = col.get(**kwargs)
|
||||||
for m in batch["metadatas"]:
|
rows = batch["metadatas"]
|
||||||
|
for m in rows:
|
||||||
r = m.get("room", "unknown")
|
r = m.get("room", "unknown")
|
||||||
rooms[r] = rooms.get(r, 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:
|
except Exception as e:
|
||||||
return {
|
return {
|
||||||
"wing": wing or "all",
|
"wing": wing or "all",
|
||||||
@@ -272,19 +281,22 @@ def tool_get_taxonomy():
|
|||||||
batch_size = 5000
|
batch_size = 5000
|
||||||
offset = 0
|
offset = 0
|
||||||
try:
|
try:
|
||||||
total = col.count()
|
col.count() # verify collection is accessible
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return {"taxonomy": {}, "error": str(e)}
|
return {"taxonomy": {}, "error": str(e)}
|
||||||
while offset < total:
|
while True:
|
||||||
try:
|
try:
|
||||||
batch = col.get(include=["metadatas"], limit=batch_size, offset=offset)
|
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")
|
w = m.get("wing", "unknown")
|
||||||
r = m.get("room", "unknown")
|
r = m.get("room", "unknown")
|
||||||
if w not in taxonomy:
|
if w not in taxonomy:
|
||||||
taxonomy[w] = {}
|
taxonomy[w] = {}
|
||||||
taxonomy[w][r] = taxonomy[w].get(r, 0) + 1
|
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:
|
except Exception as e:
|
||||||
return {
|
return {
|
||||||
"taxonomy": taxonomy,
|
"taxonomy": taxonomy,
|
||||||
|
|||||||
Reference in New Issue
Block a user