diff --git a/agents/nucliadb/src/hyperforge_nucliadb/basic_ask_agent.py b/agents/nucliadb/src/hyperforge_nucliadb/basic_ask_agent.py index ba41d79..d3e47ec 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb/basic_ask_agent.py +++ b/agents/nucliadb/src/hyperforge_nucliadb/basic_ask_agent.py @@ -1554,6 +1554,17 @@ async def inner_facets( logger.info( f"No labels selected for faceted search on question {question}, skipping." ) + context = Context( + agent_id=self.agent_id, + original_question_uuid=memory.original_question_uuid, + actual_question_uuid=None, + missing=question, + question=question, + source=source.id, + chunks=[], + agent="basic_ask_facets", + title=f"Faceted search on {source.id} Knowledge Box", + ) else: # We are assuming resources equals fields here facets, total = await ndb.field_labels(labelsets=labels_selected) @@ -1593,7 +1604,11 @@ async def inner_facets( await memory.add_step( step_module=self.config.module, step_title=self.step_title("Faceted count"), - step_reason="Faceted search performed", + step_reason=( + "Faceted search performed" + if labels_selected + else "Faceted search skipped" + ), step_value=selected_labels.get("reasoning", ""), timeit=time() - t0, input_nuclia_tokens=input_nt, diff --git a/agents/nucliadb/tests/test_facets_count.py b/agents/nucliadb/tests/test_facets_count.py new file mode 100644 index 0000000..c2e42d7 --- /dev/null +++ b/agents/nucliadb/tests/test_facets_count.py @@ -0,0 +1,72 @@ +from unittest.mock import AsyncMock, Mock + +import pytest +from hyperforge.models import Context, Source + +from hyperforge_nucliadb.basic_ask_agent import BasicAskAgent +from hyperforge_nucliadb.basic_ask_config import BasicAskAgentConfig + + +@pytest.mark.asyncio +@pytest.mark.parametrize("label_sets", [[], ["topic"]]) +async def test_facets_count_handles_label_selection(mocker, label_sets): + agent = BasicAskAgent(BasicAskAgentConfig(sources=["kb"]), agent_id="facets-agent") + source = Source.model_construct(id="kb") + mocker.patch( + "hyperforge_nucliadb.basic_ask_agent.choose_source", + new=AsyncMock(return_value=[source]), + ) + driver = Mock( + labels=AsyncMock(return_value={"topic": ["science"]}), + field_labels=AsyncMock(return_value=({"topic": {"science": 2}}, 2)), + ) + mocker.patch( + "hyperforge_nucliadb.basic_ask_agent.get_ndb_driver", return_value=driver + ) + memory = Mock(original_question_uuid="question-id", add_step=AsyncMock()) + manager = Mock( + execute_json=AsyncMock( + return_value=( + {"label_sets": label_sets, "reasoning": "Label selection"}, + 1, + 2, + ) + ), + execute=AsyncMock(return_value=("There are 2 science documents.", 3, 4, 0)), + ) + + contexts = await agent.facets_count(memory, manager, question="How many?") + + assert len(contexts) == 1 + context = contexts[0] + assert isinstance(context, Context) + assert context.agent_id == "facets-agent" + assert context.original_question_uuid == "question-id" + assert context.actual_question_uuid is None + assert context.question == "How many?" + assert context.source == "kb" + assert context.agent == "basic_ask_facets" + assert context.title == "Faceted search on kb Knowledge Box" + memory.add_step.assert_awaited_once() + step = memory.add_step.call_args.kwargs + assert step["step_value"] == "Label selection" + assert step["step_agent_path"] == "/context/facets-agent/facets_count" + + if label_sets: + assert context.missing is None + assert len(context.chunks) == 1 + assert context.chunks[0].chunk_id == "facets_count_result-0" + assert context.chunks[0].text == "There are 2 science documents." + driver.field_labels.assert_awaited_once_with(labelsets=label_sets) + manager.execute.assert_awaited_once() + assert step["step_reason"] == "Faceted search performed" + assert step["input_nuclia_tokens"] == 4 + assert step["output_nuclia_tokens"] == 6 + else: + assert context.missing == "How many?" + assert context.chunks == [] + driver.field_labels.assert_not_awaited() + manager.execute.assert_not_awaited() + assert step["step_reason"] == "Faceted search skipped" + assert step["input_nuclia_tokens"] == 1 + assert step["output_nuclia_tokens"] == 2