Set agentic client type for RAO ask and find requests

agents/nucliadb/README.md

@@ -3,4 +3,14 @@
 Basic ASK - Basic initial agent
 Advanced ASK (deprecated)
 ASK (experimental)
-Sync (Sync Service oauth linked)
+Sync (Sync Service oauth linked)
+
+## Client identification
+
+NucliaDB `/ask` and `/find` requests made by these agents use
+`X-NDB-Client: agentic`, including when invoked as tools by a smart agent in a
+RAO workflow. The Sync driver inherits the same behavior. Other request headers
+are preserved, and caller-provided headers are not modified.
+
+Workflow memory searches through `/find` also use the `agentic` client type.
+The receiving NucliaDB API must support the `agentic` client type.

agents/nucliadb/src/hyperforge_nucliadb/driver.py

@@ -303,16 +303,21 @@ async def synonyms(self, sentence: str) -> str:
                 result.append(word)
         return " ".join(result)
 
-    async def ask(self, item: AskRequest, headers={}) -> SyncAskResponse:
+    async def ask(
+        self, item: AskRequest, headers: dict[str, str] | None = None
+    ) -> SyncAskResponse:
+        request_headers = httpx.Headers(headers)
+        request_headers["X-NDB-Client"] = "agentic"
         return await self.driver.ask(
-            kbid=self.config.kbid, content=item, headers=headers
+            kbid=self.config.kbid, content=item, headers=dict(request_headers)
         )
 
     async def find(
         self, q: str, filters: List[str] = [], rids: List[str] = []
     ) -> KnowledgeboxFindResults:
         return await self.driver.find(
             kbid=self.config.kbid,
+            headers={"X-NDB-Client": "agentic"},
             content=FindRequest(
                 query=q,
                 resource_filters=rids,
@@ -326,6 +331,7 @@ async def find_raw(self, q: FindRequest) -> KnowledgeboxFindResults:
         return await self.driver.find(
             kbid=self.config.kbid,
             content=q,
+            headers={"X-NDB-Client": "agentic"},
         )
 
     async def graph_relations(

agents/nucliadb/tests/test_driver.py

@@ -1,16 +1,21 @@
+import json
 from types import SimpleNamespace
 from unittest.mock import AsyncMock
 
+import httpx
 import pytest
+from nucliadb_models.search import AskRequest, FindRequest
+from nucliadb_sdk.v2 import NucliaDBAsync
 
 from hyperforge_nucliadb.driver import NucliaDBDriver
+from hyperforge_nucliadb.sync.driver import SyncDriver
 
 
-def make_driver(client):
-    return NucliaDBDriver.model_construct(
+def make_driver(client, driver_type=NucliaDBDriver):
+    return driver_type.model_construct(
         provider="nucliadb",
         name="test",
-        config=SimpleNamespace(kbid="test-kb"),
+        config=SimpleNamespace(kbid="test-kb", filters=[]),
         driver=client,
         manager=SimpleNamespace(),
         _synonyms=None,
@@ -32,3 +37,63 @@ async def test_labels_uses_complete_labelsets_response():
 
     client.get_labelsets.assert_awaited_once_with(kbid="test-kb")
     client.get_labelset.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("driver_type", [NucliaDBDriver, SyncDriver])
+@pytest.mark.parametrize("operation", ["ask", "find", "find_raw"])
+async def test_search_requests_use_agentic_client_type(driver_type, operation):
+    def handle(request: httpx.Request) -> httpx.Response:
+        endpoint = "ask" if operation == "ask" else "find"
+        assert request.url.path.endswith(f"/kb/test-kb/{endpoint}")
+        assert request.headers.get_list("x-ndb-client") == ["agentic"]
+        assert request.headers["x-test-header"] == "preserved"
+        assert json.loads(request.content)["query"] == "question"
+        result = {"resources": {}, "total": 0}
+        if operation == "ask":
+            result = {
+                "status": "success",
+                "answer": "answer",
+                "retrieval_results": result,
+            }
+        return httpx.Response(200, json=result)
+
+    client = NucliaDBAsync(
+        url="http://nucliadb.test/api",
+        headers={"X-NDB-Client": "api", "x-test-header": "preserved"},
+        _httpx_transport=httpx.MockTransport(handle),
+    )
+    driver = make_driver(client, driver_type)
+
+    async with client.session:
+        if operation == "ask":
+            result = await driver.ask(AskRequest(query="question"))
+            assert result.answer == "answer"
+        elif operation == "find":
+            result = await driver.find("question")
+            assert result.total == 0
+        else:
+            result = await driver.find_raw(FindRequest(query="question"))
+            assert result.total == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "client_header", ["X-NDB-Client", "x-ndb-client", "X-NdB-ClIeNt"]
+)
+async def test_ask_preserves_caller_headers_without_mutating_them(client_header):
+    client = AsyncMock()
+    driver = make_driver(client)
+    headers = {client_header: "widget", "X-Synchronous": "true"}
+    original_headers = headers.copy()
+    request = AskRequest(query="question")
+
+    await driver.ask(request, headers=headers)
+
+    kwargs = client.ask.await_args.kwargs
+    sent_headers = httpx.Headers(kwargs["headers"])
+    assert sent_headers.get_list("x-ndb-client") == ["agentic"]
+    assert sent_headers["X-Synchronous"] == "true"
+    assert kwargs["kbid"] == "test-kb"
+    assert kwargs["content"] is request
+    assert headers == original_headers

hyperforge/src/hyperforge/memory/memory.py

@@ -380,7 +380,9 @@ async def search_in_questions(self, question: str, all: bool):
         )
         if not all and self.id:
             request.resource_filters = [self.id]
-        return await self.nucliadb_search.find(kbid=self.kbid, content=request)
+        return await self.nucliadb_search.find(
+            kbid=self.kbid, content=request, headers={"X-NDB-Client": "agentic"}
+        )
 
     async def qa_history(self) -> list[HistoryQuestionAnswer]:
         if self.kbid is None:

hyperforge/tests/unit/arag/test_session_memory.py

@@ -0,0 +1,42 @@
+import json
+
+import httpx
+import pytest
+from nucliadb_sdk.v2 import NucliaDBAsync
+
+from hyperforge.memory.memory import SessionMemory
+from hyperforge.models import MemoryConfig
+from hyperforge.server.cache import NoCache
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("all_sessions", [False, True])
+async def test_history_search_uses_agentic_client_type(all_sessions):
+    def handle(request: httpx.Request) -> httpx.Response:
+        assert request.url.path.endswith("/kb/test-kb/find")
+        assert request.headers.get_list("x-ndb-client") == ["agentic"]
+        assert request.headers["x-test-header"] == "preserved"
+        payload = json.loads(request.content)
+        assert payload["query"] == "previous question"
+        if all_sessions:
+            assert "resource_filters" not in payload
+        else:
+            assert payload["resource_filters"] == ["session-1"]
+        return httpx.Response(200, json={"resources": {}, "total": 0})
+
+    client = NucliaDBAsync(
+        url="http://nucliadb.test/api",
+        headers={"X-NDB-Client": "api", "x-test-header": "preserved"},
+        _httpx_transport=httpx.MockTransport(handle),
+    )
+    session = SessionMemory(
+        MemoryConfig(), agent_id="agent-1", workflow_id="default", cache=NoCache()
+    )
+    session.kbid = "test-kb"
+    session.nucliadb_search = client
+    session.init("session-1")
+
+    async with client.session:
+        result = await session.search_in_questions("previous question", all_sessions)
+
+    assert result.total == 0