Flag to store a deep copy of the request body in the response

tests/unit/test_application.py

@@ -1694,3 +1694,131 @@ def stream(*args, **kwargs):
         assert list(lines) == ["event: token", 'data: {"token":"V"}']
         assert client.closed
         client.close.assert_called_once()
+
+
+class TestQueryKeepRequestBody(unittest.TestCase):
+    """Tests for the `keep_request_body` flag on sync query methods."""
+
+    def _mock_client(self, MockClient):
+        mock_client_instance = Mock()
+        MockClient.return_value = mock_client_instance
+        mock_client_instance.close = Mock()
+
+        status_response = create_mock_httpr_response(status_code=200)
+        search_response = create_mock_httpr_response(
+            status_code=200, text="{}", url="http://localhost:8080/search/"
+        )
+        mock_client_instance.get.return_value = status_response
+        mock_client_instance.post.return_value = search_response
+        return mock_client_instance
+
+    @patch("vespa.application.httpr.Client")
+    def test_request_body_none_by_default(self, MockClient):
+        self._mock_client(MockClient)
+
+        app = Vespa(url="http://localhost", port=8080)
+        r = app.query(body={"yql": "select * from sources * where true"})
+
+        self.assertIsNone(r.request_body)
+
+    @patch("vespa.application.httpr.Client")
+    def test_request_body_stored_when_flag_true(self, MockClient):
+        self._mock_client(MockClient)
+        body = {"yql": "select * from sources * where true", "nested": {"a": 1}}
+
+        app = Vespa(url="http://localhost", port=8080)
+        r = app.query(body=body, keep_request_body=True)
+
+        self.assertEqual(r.request_body, body)
+
+    @patch("vespa.application.httpr.Client")
+    def test_request_body_is_deep_copy(self, MockClient):
+        self._mock_client(MockClient)
+        body = {"yql": "select * from sources * where true", "nested": {"a": 1}}
+
+        app = Vespa(url="http://localhost", port=8080)
+        r = app.query(body=body, keep_request_body=True)
+
+        self.assertIsNot(r.request_body, body)
+        self.assertIsNot(r.request_body["nested"], body["nested"])
+
+        # Mutating the original body after the call must not affect the stored copy.
+        body["nested"]["a"] = 999
+        self.assertEqual(r.request_body["nested"]["a"], 1)
+
+    @patch("vespa.application.httpr.Client")
+    def test_keep_request_body_not_sent_as_query_param(self, MockClient):
+        mock_client_instance = self._mock_client(MockClient)
+
+        app = Vespa(url="http://localhost", port=8080)
+        app.query(
+            body={"yql": "select * from sources * where true"},
+            keep_request_body=True,
+        )
+
+        _, call_kwargs = mock_client_instance.post.call_args
+        self.assertNotIn("keep_request_body", call_kwargs.get("params", {}))
+
+    @patch("vespa.application.httpr.Client")
+    def test_request_body_none_via_session(self, MockClient):
+        self._mock_client(MockClient)
+
+        app = Vespa(url="http://localhost", port=8080)
+        with app.syncio() as session:
+            r = session.query(body={"yql": "select * from sources * where true"})
+
+        self.assertIsNone(r.request_body)
+
+    @patch("vespa.application.httpr.Client")
+    def test_request_body_stored_via_session(self, MockClient):
+        self._mock_client(MockClient)
+        body = {"yql": "select * from sources * where true"}
+
+        app = Vespa(url="http://localhost", port=8080)
+        with app.syncio() as session:
+            r = session.query(body=body, keep_request_body=True)
+
+        self.assertEqual(r.request_body, body)
+        self.assertIsNot(r.request_body, body)
+
+
+@pytest.mark.asyncio
+class TestAsyncQueryKeepRequestBody:
+    """Tests for the `keep_request_body` flag on `VespaAsync.query`."""
+
+    async def test_request_body_none_by_default(self):
+        app = Vespa(url="http://localhost", port=8080)
+        vespa_async = VespaAsync(app)
+        vespa_async._make_request = AsyncMock(
+            return_value=create_mock_httpr_response(
+                status_code=200,
+                json_data={"root": {}},
+                url="http://localhost:8080/search/",
+            )
+        )
+
+        r = await vespa_async.query(body={"yql": "select * from sources * where true"})
+
+        assert r.request_body is None
+
+    async def test_request_body_stored_when_flag_true(self):
+        app = Vespa(url="http://localhost", port=8080)
+        vespa_async = VespaAsync(app)
+        vespa_async._make_request = AsyncMock(
+            return_value=create_mock_httpr_response(
+                status_code=200,
+                json_data={"root": {}},
+                url="http://localhost:8080/search/",
+            )
+        )
+        body = {"yql": "select * from sources * where true", "nested": {"a": 1}}
+
+        r = await vespa_async.query(body=body, keep_request_body=True)
+
+        assert r.request_body == body
+        assert r.request_body is not body
+        assert r.request_body["nested"] is not body["nested"]
+
+        # Mutating the original body after the call must not affect the stored copy.
+        body["nested"]["a"] = 999
+        assert r.request_body["nested"]["a"] == 1

tests/unit/test_io.py

@@ -1,7 +1,7 @@
 # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 import unittest
-from vespa.io import VespaVisitResponse, VespaQueryResponse
+from vespa.io import VespaVisitResponse, VespaQueryResponse, VespaResponse
 
 
 class TestVespaVisitResult(unittest.TestCase):
@@ -158,3 +158,66 @@ def test_hits(self):
                 }
             ],
         )
+
+
+class TestVespaQueryResponseEquality(unittest.TestCase):
+    def setUp(self) -> None:
+        self.json = {"root": {"fields": {"totalCount": 1}}}
+        self.status_code = 200
+        self.url = "http://localhost:8080/search/"
+        self.body = {"yql": "select * from sources * where true"}
+
+    def _make(self, request_body=None):
+        return VespaQueryResponse(
+            json=self.json,
+            status_code=self.status_code,
+            url=self.url,
+            request_body=request_body,
+        )
+
+    def test_equal_when_request_body_matches(self):
+        self.assertEqual(self._make(), self._make())
+        self.assertEqual(
+            self._make(request_body=self.body), self._make(request_body=self.body)
+        )
+
+    def test_equal_uses_value_equality_not_identity(self):
+        # Equal but distinct dict objects (e.g. one is a deep copy of the other).
+        a = self._make(request_body={"yql": "x", "nested": {"n": 1}})
+        b = self._make(request_body={"yql": "x", "nested": {"n": 1}})
+        self.assertIsNot(a.request_body, b.request_body)
+        self.assertEqual(a, b)
+
+    def test_not_equal_when_request_body_differs(self):
+        a = self._make(request_body={"yql": "select * from sources * where a"})
+        b = self._make(request_body={"yql": "select * from sources * where b"})
+        self.assertNotEqual(a, b)
+
+    def test_not_equal_when_one_request_body_is_none(self):
+        a = self._make(request_body=None)
+        b = self._make(request_body=self.body)
+        self.assertNotEqual(a, b)
+        self.assertNotEqual(b, a)
+
+    def test_not_equal_to_unrelated_type(self):
+        a = self._make(request_body=self.body)
+        self.assertNotEqual(a, "not a response")
+        self.assertNotEqual(
+            a,
+            VespaResponse(
+                json=self.json,
+                status_code=self.status_code,
+                url=self.url,
+                operation_type="query",
+            ),
+        )
+
+    def test_base_fields_still_compared(self):
+        a = self._make(request_body=self.body)
+        b = VespaQueryResponse(
+            json={"root": {"fields": {"totalCount": 2}}},
+            status_code=self.status_code,
+            url=self.url,
+            request_body=self.body,
+        )
+        self.assertNotEqual(a, b)

vespa/application.py

@@ -2,6 +2,7 @@
 
 import sys
 import asyncio
+import copy
 import traceback
 import concurrent.futures
 import warnings
@@ -573,6 +574,7 @@ def query(
         groupname: Optional[str] = None,
         streaming: bool = False,
         profile: bool = False,
+        keep_request_body: bool = False,
         **kwargs,
     ) -> Union[VespaQueryResponse, Generator[str, None, None]]:
         """
@@ -585,6 +587,8 @@ def query(
             groupname (str, optional): The groupname used with streaming search.
             streaming (bool, optional): Whether to use streaming mode (SSE). Defaults to False.
             profile (bool, optional): Add profiling parameters to the query (response may be large). Defaults to False.
+            keep_request_body (bool, optional): If True, store a deep copy of `body` on the returned
+                response as `response.request_body`. Defaults to False. Has no effect when streaming=True.
             **kwargs (dict, optional): Extra Vespa Query API parameters.
 
         Returns:
@@ -605,6 +609,7 @@ def query(
                 groupname=groupname,
                 streaming=False,
                 profile=profile,
+                keep_request_body=keep_request_body,
                 **kwargs,
             )
 
@@ -1737,6 +1742,7 @@ def query(
         groupname: Optional[str] = None,
         streaming: bool = False,
         profile: bool = False,
+        keep_request_body: bool = False,
         **kwargs,
     ) -> Union[VespaQueryResponse, Generator[str, None, None]]:
         """
@@ -1747,6 +1753,8 @@ def query(
             groupname (str, optional): The groupname used in streaming search.
             streaming (bool, optional): Whether to use streaming mode (SSE). Defaults to False.
             profile (bool, optional): Add profiling parameters to the query (response may be large). Defaults to False.
+            keep_request_body (bool, optional): If True, store a deep copy of `body` on the returned
+                response as `response.request_body`. Defaults to False. Has no effect when streaming=True.
             **kwargs (dict, optional): Additional valid Vespa HTTP Query API parameters. See: <https://docs.vespa.ai/en/reference/query-api-reference.html>.
 
         Returns:
@@ -1776,6 +1784,7 @@ def query(
                 json=response.json(),
                 status_code=response.status_code,
                 url=str(response.url),
+                request_body=copy.deepcopy(body) if keep_request_body else None,
             )
 
     def _query_streaming(
@@ -2325,6 +2334,7 @@ async def query(
         groupname: Optional[str] = None,
         profile: bool = False,
         retry_policy: Optional[AsyncRetrying] = None,
+        keep_request_body: bool = False,
         **kwargs,
     ) -> VespaQueryResponse:
         """
@@ -2335,6 +2345,8 @@ async def query(
             groupname (str, optional): The groupname used in streaming search.
             profile (bool, optional): Add profiling parameters to the query (response may be large). Defaults to False.
             retry_policy (AsyncRetrying, optional): Custom tenacity retry policy. Defaults to ``vespa.retries.QUERY_RETRY`` (five attempts with random exponential wait time).
+            keep_request_body (bool, optional): If True, store a deep copy of `body` on the returned
+                response as `response.request_body`. Defaults to False.
             **kwargs (dict, optional): Additional valid Vespa HTTP Query API parameters.
 
         Returns:
@@ -2362,6 +2374,7 @@ async def _do_query() -> VespaQueryResponse:
                 json=_response_json(response),
                 status_code=response.status_code,
                 url=str(response.url),
+                request_body=copy.deepcopy(body) if keep_request_body else None,
             )
 
         return await retry_policy(_do_query)

vespa/io.py

@@ -16,7 +16,10 @@ def __init__(
         self.operation_type = operation_type
 
     def __eq__(self, other: object) -> bool:
-        if not isinstance(other, self.__class__):
+        # Compare exact types (not isinstance) so equality stays symmetric
+        # across subclasses, e.g. a VespaQueryResponse is never equal to a
+        # plain VespaResponse even if their base fields happen to match.
+        if type(other) is not type(self):
             return NotImplemented
         return (
             self.json == other.json
@@ -53,8 +56,16 @@ def __init__(self, json, status_code, url, request_body=None) -> None:
         )
         self._request_body = request_body
 
+    def __eq__(self, other: object) -> bool:
+        base_eq = super().__eq__(other)
+        if base_eq is NotImplemented:
+            return NotImplemented
+        return base_eq and self.request_body == other.request_body
+
     @property
     def request_body(self) -> Optional[Dict]:
+        """A deep copy of the request body, set only when the query was made with
+        `keep_request_body=True`. `None` otherwise."""
         return self._request_body
 
     @property