diff --git a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py index 3ad3a5db16..acd2859b72 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py @@ -69,6 +69,32 @@ def prompt_output(request): status=status.HTTP_400_BAD_REQUEST, ) + # Both of these are indexed by key inside OutputManagerHelper. + # handle_prompt_output_update: `metadata` at its lines 135-139 (context, + # challenge_data, highlight_data, confidence_data, word_confidence_data, + # all unconditional and reached before the `if not prompts` early exit), + # and `outputs` at line 179. A non-mapping in either raises AttributeError + # deep in the helper and surfaces as an opaque 500. + # + # Reachable from ordinary traffic, not just a malformed client: single-pass + # extraction passes the LLM's parsed JSON straight through as `outputs`, + # and the repair returns a list whenever the model adds commentary or + # answers in prose. Rejecting both here with a reason means no executor can + # 500 the backend this way (UN-4017). + for field_name, value in (("outputs", outputs), ("metadata", metadata)): + if not isinstance(value, dict): + return JsonResponse( + { + "success": False, + "error": ( + f"{field_name} must be a JSON object, got " + f"{type(value).__name__}. A JSON array is valid JSON " + "but cannot be indexed by key." + ), + }, + status=status.HTTP_400_BAD_REQUEST, + ) + try: from prompt_studio.prompt_studio_output_manager_v2.output_manager_helper import ( OutputManagerHelper, diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 6d8763e05a..ba157801ce 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -1875,14 +1875,52 @@ def _handle_response( "message": IndexingStatus.DOCUMENT_BEING_INDEXED.value, } + outputs = response["output"] + metadata = response["metadata"] + # Same guard as the internal API (UN-4017). This is the in-backend + # execution path — it dispatches the identical single_pass_extraction + # executor, so it can receive the identical bad shape. Without this, + # handle_prompt_output_update does outputs.get(prompt.prompt_key) on a + # list and raises AttributeError, which surfaces as a bare 500. + # `metadata` is indexed five times at the top of + # handle_prompt_output_update, unconditionally and before its + # `if not prompts` early exit, so it has the same exposure as `outputs` + # and had been left unchecked here. + for field_name, value in (("outputs", outputs), ("metadata", metadata)): + if isinstance(value, dict): + continue + # Name the type rather than asserting a shape: this fires for + # NoneType, int and bool too, and "LLM returned a JSON array + # (got NoneType)" contradicts itself. + detail = ( + f"LLM response could not be used as the {field_name} map — a " + "single JSON object keyed by field name is expected (got " + f"{type(value).__name__})." + ) + # Only for `outputs`. `metadata` is assembled by the executor, not + # returned by the LLM, so telling the user to rephrase a prompt + # would be a dead end for a defect that is ours — the same + # misdirection the shape claim above was rewritten to avoid. + if is_single_pass and field_name == "outputs": + detail += ( + " In single-pass extraction all prompts share one response," + " so a prompt that asks for a list or for separate JSON" + " entries can change the shape of the entire result." + " Rephrase that prompt to describe the value of its own" + " field, or run these prompts with single-pass extraction" + " turned off." + ) + logger.error("%s run_id=%s document_id=%s", detail, run_id, document_id) + raise AnswerFetchError(detail, status_code=422) + return OutputManagerHelper.handle_prompt_output_update( run_id=run_id, prompts=prompts, - outputs=response["output"], + outputs=outputs, document_id=document_id, is_single_pass_extract=is_single_pass, profile_manager_id=profile_manager_id, - metadata=response["metadata"], + metadata=metadata, ) @staticmethod diff --git a/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py b/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py new file mode 100644 index 0000000000..34f7dbea7a --- /dev/null +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py @@ -0,0 +1,215 @@ +"""UN-4017: prompt_output rejects a non-mapping ``outputs`` at the boundary. + +``outputs`` is indexed by prompt key downstream — OutputManagerHelper. +handle_prompt_output_update does ``outputs.get(prompt.prompt_key)`` — so a list +raised AttributeError deep inside the helper and surfaced to the caller as a +bare 500 with no usable reason. + +This is reachable from real traffic, not just a malformed client: single-pass +extraction passes the LLM's parsed JSON straight through as the outputs map, +and that parse returns a list whenever the model wraps its answer in prose, a +reasoning block, or a stray fence marker. Validating here means no future +executor can 500 the backend the same way. +""" + +import json +from unittest.mock import MagicMock, patch + +from prompt_studio.prompt_studio_core_v2.internal_views import prompt_output + +_VIEWS = "prompt_studio.prompt_studio_core_v2.internal_views" + + +def _request(outputs): + request = MagicMock() + # require_http_methods reads request.method; a MagicMock attribute is not + # the string "POST" and the view short-circuits with 405. + request.method = "POST" + request.body = json.dumps( + { + "run_id": "run-1", + "prompt_ids": ["p1"], + "outputs": outputs, + "document_id": "doc-1", + "is_single_pass_extract": True, + } + ) + return request + + +def _body(response): + return json.loads(response.content) + + +def test_list_outputs_rejected_with_400_and_a_reason(): + """The regression: this used to reach the helper and raise AttributeError.""" + response = prompt_output(_request([{"invoice_number": "INV-001"}, {"b": 2}])) + assert response.status_code == 400 + body = _body(response) + assert body["success"] is False + assert "outputs must be a JSON object" in body["error"] + assert "list" in body["error"] + # A JSON array is valid JSON — the message must not imply it was malformed. + assert "malformed" not in body["error"].lower() + assert "invalid json" not in body["error"].lower() + + +def test_non_mapping_metadata_rejected(): + """metadata is indexed five times at helper lines 135-139, before the + `if not prompts` early exit — so it 500s even with valid outputs.""" + request = MagicMock() + request.method = "POST" + request.body = json.dumps( + { + "prompt_ids": ["p1"], + "document_id": "doc-1", + "outputs": {"invoice_number": "INV-001"}, + "metadata": [], + } + ) + response = prompt_output(request) + assert response.status_code == 400 + assert "metadata must be a JSON object" in _body(response)["error"] + + +def test_non_mapping_outputs_rejected(): + for outputs in ("a string", 42, True, [], [1, 2]): + response = prompt_output(_request(outputs)) + assert response.status_code == 400, f"{outputs!r} was accepted" + + +def test_helper_is_never_reached_for_invalid_outputs(): + """Rejected at the boundary, before any ORM or helper work happens. + + The ORM is patched so the assertion cannot pass by accident: without it, + removing the guard makes `filter()` raise on the fake prompt id and the + helper goes uncalled for the wrong reason. + """ + with patch( + "prompt_studio.prompt_studio_v2.models.ToolStudioPrompt.objects" + ) as prompts, patch( + "prompt_studio.prompt_studio_output_manager_v2." + "output_manager_helper.OutputManagerHelper.handle_prompt_output_update" + ) as handler: + prompts.filter.return_value.order_by.return_value = [] + response = prompt_output(_request([{"a": 1}])) + handler.assert_not_called() + assert response.status_code == 400 + + +def test_dict_outputs_still_reach_the_helper(): + """The guard must not reject the ordinary case.""" + with patch( + "prompt_studio.prompt_studio_v2.models.ToolStudioPrompt.objects" + ) as prompts, patch( + "prompt_studio.prompt_studio_output_manager_v2." + "output_manager_helper.OutputManagerHelper.handle_prompt_output_update" + ) as handler: + prompts.filter.return_value.order_by.return_value = [] + handler.return_value = [] + response = prompt_output(_request({"invoice_number": "INV-001"})) + handler.assert_called_once() + assert response.status_code == 200 + + +def test_missing_outputs_defaults_to_empty_dict_and_is_accepted(): + """Absent `outputs` has always meant {}; the guard must not change that.""" + request = MagicMock() + request.method = "POST" + request.body = json.dumps({"prompt_ids": ["p1"], "document_id": "doc-1"}) + with patch( + "prompt_studio.prompt_studio_v2.models.ToolStudioPrompt.objects" + ) as prompts, patch( + "prompt_studio.prompt_studio_output_manager_v2." + "output_manager_helper.OutputManagerHelper.handle_prompt_output_update" + ) as handler: + prompts.filter.return_value.order_by.return_value = [] + handler.return_value = [] + response = prompt_output(request) + assert response.status_code == 200 + + +# --- The in-backend execution path ------------------------------------- +# +# prompt_studio_helper._handle_response is the other route to +# handle_prompt_output_update. It dispatches the same single_pass_extraction +# executor, so it receives the same shapes; guarding only the internal API +# would have left this path still able to 500. + +from prompt_studio.prompt_studio_core_v2.exceptions import AnswerFetchError # noqa: E402 +from prompt_studio.prompt_studio_core_v2.prompt_studio_helper import ( # noqa: E402 + PromptStudioHelper, +) + + +def _handle(outputs, is_single_pass=True): + return PromptStudioHelper._handle_response( + response={"output": outputs, "metadata": {}, "status": "COMPLETED"}, + run_id="run-1", + prompts=[], + document_id="doc-1", + is_single_pass=is_single_pass, + ) + + +def test_in_backend_path_rejects_a_list_with_422(): + try: + _handle([{"invoice_number": "INV-001"}, {"b": 2}]) + except AnswerFetchError as exc: + assert exc.status_code == 422 + assert "outputs map" in str(exc.detail) + else: + raise AssertionError("a list was accepted on the in-backend path") + + +def test_in_backend_single_pass_message_points_at_the_prompt(): + try: + _handle([{"a": 1}], is_single_pass=True) + except AnswerFetchError as exc: + assert "all prompts share one response" in str(exc.detail) + else: + raise AssertionError("expected AnswerFetchError") + + +def test_in_backend_single_prompt_omits_the_single_pass_advice(): + """That advice is only true of single pass; it would misdirect otherwise.""" + try: + _handle([{"a": 1}], is_single_pass=False) + except AnswerFetchError as exc: + assert "outputs map" in str(exc.detail) + assert "all prompts share one response" not in str(exc.detail) + else: + raise AssertionError("expected AnswerFetchError") + + +def test_in_backend_path_still_accepts_a_dict(): + with patch( + "prompt_studio.prompt_studio_output_manager_v2." + "output_manager_helper.OutputManagerHelper.handle_prompt_output_update" + ) as handler: + handler.return_value = [] + _handle({"invoice_number": "INV-001"}) + handler.assert_called_once() + + +def test_metadata_rejection_does_not_tell_the_user_to_rephrase_a_prompt(): + """`metadata` is executor-assembled, never LLM output. + + Rewording a prompt cannot change it, so the single-pass advice would send + the user down a dead end for a defect that is ours. + """ + try: + PromptStudioHelper._handle_response( + response={"output": {"a": 1}, "metadata": [], "status": "COMPLETED"}, + run_id="run-1", + prompts=[], + document_id="doc-1", + is_single_pass=True, + ) + except AnswerFetchError as exc: + detail = str(exc.detail) + assert "metadata map" in detail + assert "all prompts share one response" not in detail + assert "Rephrase" not in detail + else: + raise AssertionError("non-dict metadata was accepted")