diff --git a/pageindex/agent.py b/pageindex/agent.py index 54c104b..e5de60d 100644 --- a/pageindex/agent.py +++ b/pageindex/agent.py @@ -20,7 +20,7 @@ OPEN_SYSTEM_PROMPT = """ You are PageIndex, a document QA assistant. TOOL USE: - Call list_documents() to see available documents; use doc_name and doc_description to pick which doc(s) are relevant. -- Call get_document(doc_id) to confirm status and page/line count. +- Call get_document(doc_id) to confirm the document's name and type. - Call get_document_structure(doc_id) to identify relevant page ranges. - Call get_page_content(doc_id, pages="5-7") with tight ranges; never fetch the whole document. - Before each tool call, output one short sentence explaining the reason. @@ -33,7 +33,7 @@ Answer based only on tool output. Be concise. SCOPED_SYSTEM_PROMPT = """ You are PageIndex, a document QA assistant. TOOL USE: -- Call get_document(doc_id) to confirm status and page/line count. +- Call get_document(doc_id) to confirm the document's name and type. - Call get_document_structure(doc_id) to identify relevant page ranges. - Call get_page_content(doc_id, pages="5-7") with tight ranges; never fetch the whole document. - Before each tool call, output one short sentence explaining the reason. diff --git a/tests/test_agent.py b/tests/test_agent.py index 6ec5b4d..162ef9a 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -20,6 +20,16 @@ def test_scoped_prompt_omits_list_documents(): assert "get_page_content" in SCOPED_SYSTEM_PROMPT +def test_prompts_get_document_guidance_matches_returned_fields(): + # Regression: get_document returns doc_name/doc_type/doc_description — neither + # backend returns a page/line count, and the local backend has no status + # field. The prompt must not send the agent hunting for fields that don't + # exist (degrades QA), so it references only name and type (like the demo). + for prompt in (OPEN_SYSTEM_PROMPT, SCOPED_SYSTEM_PROMPT): + assert "page/line count" not in prompt + assert "get_document(doc_id) to confirm the document's name and type" in prompt + + def test_wrap_with_doc_context_cannot_be_escaped_by_untrusted_content(): """doc_name/doc_description are untrusted (doc_name is an unsanitized filename; doc_description is LLM-generated from document content). Neither