diff --git a/agents/nucliadb/README.md b/agents/nucliadb/README.md index 65c51310..08db5da4 100644 --- a/agents/nucliadb/README.md +++ b/agents/nucliadb/README.md @@ -3,4 +3,14 @@ Basic ASK - Basic initial agent Advanced ASK (deprecated) ASK (experimental) -Sync (Sync Service oauth linked) \ No newline at end of file +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. \ No newline at end of file diff --git a/agents/nucliadb/src/hyperforge_nucliadb/driver.py b/agents/nucliadb/src/hyperforge_nucliadb/driver.py index 762d08aa..2c5dc195 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/driver.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/driver.py @@ -303,9 +303,13 @@ 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( @@ -313,6 +317,7 @@ async def find( ) -> 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( diff --git a/agents/nucliadb/tests/test_driver.py b/agents/nucliadb/tests/test_driver.py index 9282dafe..42af53d2 100644 --- a/agents/nucliadb/tests/test_driver.py +++ b/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 diff --git a/hyperforge/src/hyperforge/memory/memory.py b/hyperforge/src/hyperforge/memory/memory.py index 6d1bdbf7..cfb2b3d3 100644 --- a/hyperforge/src/hyperforge/memory/memory.py +++ b/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: diff --git a/hyperforge/tests/unit/arag/test_session_memory.py b/hyperforge/tests/unit/arag/test_session_memory.py new file mode 100644 index 00000000..4f85bf45 --- /dev/null +++ b/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