Update app.py

app.py

@@ -154,15 +154,7 @@ def search(query: str, k: int = 5) -> List[int]:
     top_k_indices = scores[0].topk(k).indices.tolist()
     print("[search]", query, top_k_indices)
 
-    # Neighbor expansion for context
-    base = set(top_k_indices)
-    expanded = set(base)
-    for i in base:
-        expanded.add(i - 1)
-        expanded.add(i + 1)
-    expanded = {i for i in expanded if 0 <= i < len(images)}
-
-    return sorted(expanded)
+    return top_k_indices
 
 
 def _build_image_parts_from_indices(indices: List[int]) -> List[Dict[str, Any]]:
@@ -189,7 +181,7 @@ Act iteratively:
   1) Split the user question into 1–4 focused sub-queries. Subqueries should be asked as natural language questions in the english language, not just keywords.
   2) For each sub-query, call mcp_test_search (k=5 by default; increase to up to 10 if you need to go deep).
   3) You will receive the output of mcp_test_search as a list of indices corresponding to page numbers. Print them out and stop generating. You will be fed the corresponding pages as images in a follow-up message.
-  3) Stop early when confident; otherwise refine and repeat, running new searches. Up to 4 iterations and 20 searches in total. If info is missing, try to continue searching using new keywords and queries.
+  3) Stop early when confident; otherwise refine and repeat, running new searches. Up to 5 iterations and 20 searches in total. If info is missing, try to continue searching using new keywords and queries.
 
 Workflow:
   • Use ONLY the provided images for grounding and cite as (p.<page>).
@@ -369,7 +361,7 @@ def stream_agent(question: str,
             return
 
     # Controller: iterate rounds; if the model searched, attach those pages next
-    max_rounds = 3
+    max_rounds = 5
     round_idx = 1
     pending_indices = list(seed_indices)
 
@@ -381,7 +373,14 @@ def stream_agent(question: str,
         # If the model returned indices via the tool, use them in a fresh call
         next_indices = round_state.get("last_search_indices") or []
         if next_indices:
-            pending_indices = next_indices
+            # Neighbor expansion for context
+            base = set(next_indices)
+            expanded = set(base)
+            for i in base:
+                expanded.add(i - 1)
+                expanded.add(i + 1)
+            expanded = {i for i in expanded if 0 <= i < len(images)}
+            pending_indices = sorted(expanded)
             round_idx += 1
             continue