Update app.py

app.py

@@ -133,9 +133,21 @@ def index_from_url(url: str) -> Tuple[str, str]:
 
 def search(query: str, k: int = 5) -> List[int]:
     """
-    Search within an indexed PDF and return ONLY the indices of the most relevant pages (0-based).
+    Search within a PDF document for the most relevant pages to answer a query and return the page indexes as a list.
+    MCP tool description:
+      - name: mcp_test_search
+      - description: Search within a PDF document for the most relevant pages to answer a query.
+      - input_schema:
+          type: object
+          properties:
+            query: {type: string, description: "User query in natural language."}
+            k: {type: integer, minimum: 1, maximum: 10, default: 5. description: "Number of top pages to retrieve."}
+          required: ["query"]
+    Args:
+        query (str): Natural-language question to search for.
+        k (int): Number of top results to return (1–10).
     Returns:
-      List[int]: Sorted unique 0-based indices of pages to inspect (includes neighbor expansion).
+        indices (List[int]): Indices of the k most relevant pages
     """
     global ds, images
 
@@ -180,15 +192,15 @@ You are a PDF research agent with a single tool: mcp_test_search(query: string,
 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 5 iterations and 20 searches in total. If info is missing, try to continue searching using new keywords and queries.
+  3) You will receive the output of mcp_test_search as a list of indices corresponding to page numbers. Stop generating once all the tool calls end. You will later be fed the corresponding pages as images in a follow-up message.
+  4) Stop early when confident; otherwise refine and repeat, running new search calls when need be. Use 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>).
   • If an answer is not present, say “Not found in the provided pages.”
 
 Deliverable:
-  • Return a clear, standalone Markdown answer in the user's language. Include concise tables for lists of dates/items.
+  • Return a clear, standalone Markdown answer in the user's language. Include concise tables for lists of dates/items when useful, and cite the page numbers used for each fact.
 """
 ).strip()