From 5eb0a306ede2198e60cb69303ad87ae5f70bb395 Mon Sep 17 00:00:00 2001 From: Carmen Iniesta Date: Thu, 12 Nov 2026 10:32:45 +0100 Subject: [PATCH] use kb default --- .../src/hyperforge_nucliadb_agentic/agent.py | 36 +++- .../src/hyperforge_nucliadb_agentic/config.py | 8 +- agents/nucliadb/tests/test_agent.py | 183 +++++++++++++++++- agents/nucliadb/tests/test_config.py | 7 +- tests/db/test_transform.py | 1 + 5 files changed, 226 insertions(+), 9 deletions(-) diff --git a/agents/nucliadb/src/hyperforge_nucliadb_agentic/agent.py b/agents/nucliadb/src/hyperforge_nucliadb_agentic/agent.py index a525779..6bd76f3 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb_agentic/agent.py +++ b/agents/nucliadb/src/hyperforge_nucliadb_agentic/agent.py @@ -9,6 +9,7 @@ from hyperforge.configure import agent from hyperforge.context.agent import ContextAgent from hyperforge.definition import FunctionDefinition +from hyperforge.llm_config import LLMConfig from hyperforge.manager import Manager, build_reasoning from hyperforge.memory import Chunk, Context, QuestionMemory, Source from hyperforge.models import JSONObject @@ -67,9 +68,20 @@ async def choose_sources( question: str, ident: str, step_title: str, + model: LLMConfig | None = None, ) -> list[Source]: if len(sources) == 1: return [Source.model_construct(id=sources[0])] + kbid = None + if model is None and sources: + kbids = {get_ndb_driver(manager, source).config.kbid for source in sources} + if len(kbids) != 1: + raise ValueError( + "An explicit routing model is required for sources from different KBs" + ) + kbid = kbids.pop() + if not kbid or not kbid.strip(): + raise ValueError("A KB ID is required to inherit its generative model") return await choose_source( memory, manager, @@ -77,6 +89,8 @@ async def choose_sources( question, ident=ident, step_title=step_title, + choose_source_model=model, + kbid=kbid, ) @@ -775,6 +789,7 @@ async def ask_agent( question, ident=self.agent_id, step_title=self.step_title("Choose sources"), + model=self.config.generative_model, ) contexts: list[Context] = await asyncio.gather( @@ -812,6 +827,7 @@ async def _get_question_context( question, ident=self.agent_id, step_title=self.step_title("Choose sources"), + model=self.config.generative_model, ) missing: list[tuple[str, str] | None] = await asyncio.gather( @@ -979,6 +995,8 @@ async def prepare_ask_request( rag_strategies: list[RagStrategies], ) -> AskRequest: """Apply agent defaults, search config, public ASK overrides, and runtime query.""" + if not nucliadb_driver.config.kbid or not nucliadb_driver.config.kbid.strip(): + raise ValueError("A KB ID is required for KB source generation") # Load explicit options from the public ASK request. ask_request: AskRequest | None = None if ask_request_json: @@ -1007,14 +1025,15 @@ async def prepare_ask_request( ) # Agent defaults fill any remaining fields. - fallback_values = { + fallback_values: dict[str, object] = { "show": [ResourceProperties.BASIC, ResourceProperties.ORIGIN], "citations": CitationsType.LLM_FOOTNOTES, - "generative_model": self.config.generative_model.model_id, "rag_strategies": rag_strategies, "generate_answer": self.config.generate_inner_answer, - "reasoning": build_reasoning(self.config.generative_model), } + if self.config.generative_model is not None: + fallback_values["generative_model"] = self.config.generative_model.model_id + fallback_values["reasoning"] = build_reasoning(self.config.generative_model) return ask_request.model_copy( update={ field: value @@ -1072,6 +1091,8 @@ async def inner_rag( memory.arguments.get("ask_request"), rag_strategies, ) + if ask_request.generative_model is None: + manager.require_kb_model_scope(nucliadb_driver.config.kbid) filter_expression = await self.build_filter_expression( nucliadb_driver, @@ -1233,6 +1254,7 @@ async def facets_search( question, ident=self.agent_id, step_title=self.step_title("Choose sources"), + model=self.config.generative_model, ) contexts: list[Context] = await asyncio.gather( *[ @@ -1287,6 +1309,7 @@ async def inner_facets_search( "additionalProperties": False, }, model=self.config.generative_model, + kbid=ndb.config.kbid, tracking=memory.get_tracking_info(), ) labelsets_selected = selected_labels.get("label_sets", []) @@ -1341,6 +1364,7 @@ async def inner_facets_search( prompt, user_id="rao-facets-answer-generation", model=self.config.generative_model, + kbid=ndb.config.kbid, tracking=memory.get_tracking_info(), ) # TODO: Error if code is not success @@ -1392,6 +1416,7 @@ async def facets_count( question, ident=self.agent_id, step_title=self.step_title("Choose sources"), + model=self.config.generative_model, ) contexts: list[Context] = await asyncio.gather( *[ @@ -1441,6 +1466,7 @@ async def inner_facets( "additionalProperties": False, }, model=self.config.generative_model, + kbid=ndb.config.kbid, tracking=memory.get_tracking_info(), ) labels_selected = selected_labels.get("label_sets", []) @@ -1462,6 +1488,7 @@ async def inner_facets( prompt, user_id="rao-facets-answer-generation", model=self.config.generative_model, + kbid=ndb.config.kbid, tracking=memory.get_tracking_info(), ) # TODO: Error if code is not success @@ -1512,6 +1539,7 @@ async def catalog_search( question, ident=self.agent_id, step_title=self.step_title("Choose sources"), + model=self.config.generative_model, ) contexts: list[Context] = await asyncio.gather( *[ @@ -1557,6 +1585,7 @@ async def inner_catalog_search( "required": ["filters"], }, model=self.config.generative_model, + kbid=ndb.config.kbid, tracking=memory.get_tracking_info(), ) filters = self.parse_selected_filters(question, source, selected_filters) @@ -1587,6 +1616,7 @@ async def inner_catalog_search( prompt, user_id="rao-catalog-search-answer", model=self.config.generative_model, + kbid=ndb.config.kbid, tracking=memory.get_tracking_info(), ) # TODO: Error if code is not success diff --git a/agents/nucliadb/src/hyperforge_nucliadb_agentic/config.py b/agents/nucliadb/src/hyperforge_nucliadb_agentic/config.py index c19b330..b9e99f0 100644 --- a/agents/nucliadb/src/hyperforge_nucliadb_agentic/config.py +++ b/agents/nucliadb/src/hyperforge_nucliadb_agentic/config.py @@ -1,7 +1,7 @@ from typing import List, Literal, Optional, Tuple from hyperforge.context.agent import ContextAgentConfig -from hyperforge.llm_config import LLMConfig, LLMField, llm_defaults +from hyperforge.llm_config import LLMField from hyperforge.utils import WidgetType from pydantic import Field from pydantic.config import ConfigDict @@ -20,10 +20,10 @@ class NucliaDBAgentConfig(ContextAgentConfig): "show_in_node": True, }, ) - generative_model: LLMField = Field( - default=LLMConfig(model_id=llm_defaults.default), + generative_model: LLMField | None = Field( + default=None, title="Generative model", - description="Model used to generate answers", + description="Model used by KB source tools. If unset, inherit the Knowledge Box's configured model.", ) published_functions: Optional[Tuple[str, ...]] = Field( default=( diff --git a/agents/nucliadb/tests/test_agent.py b/agents/nucliadb/tests/test_agent.py index b256f62..955fd2f 100644 --- a/agents/nucliadb/tests/test_agent.py +++ b/agents/nucliadb/tests/test_agent.py @@ -9,6 +9,8 @@ import pytest from hyperforge.llm_config import LLMConfig +from hyperforge.memory import Source + from hyperforge_nucliadb_agentic.agent import ( NucliaDBAgent, choose_sources, @@ -26,6 +28,36 @@ class TestChooseSources: + async def test_router_forwards_inherited_model_and_common_kb( + self, mock_memory, mock_manager, mock_nucliadb_driver + ): + mock_manager.drivers.update( + {"first": mock_nucliadb_driver, "second": mock_nucliadb_driver} + ) + mock_memory.get_session_source = AsyncMock( + side_effect=[ + Source.model_construct(id="first"), + Source.model_construct(id="second"), + ] + ) + mock_memory.get_rules.return_value = [] + mock_manager.execute_json.return_value = ( + {"sources": ["first"], "reason": "Relevant source"}, + 1, + 1, + ) + sources = await choose_sources( + mock_memory, + mock_manager, + ["first", "second"], + "question", + ident="agent", + step_title="Choose sources", + ) + assert [source.id for source in sources] == ["first"] + assert mock_manager.execute_json.await_args.kwargs["model"] is None + assert mock_manager.execute_json.await_args.kwargs["kbid"] == "test-kbid" + async def test_single_source_does_not_load_routing_metadata( self, mock_memory, mock_manager ): @@ -45,7 +77,15 @@ async def test_single_source_does_not_load_routing_metadata( assert [source.id for source in sources] == ["local-kb"] choose_source.assert_not_awaited() - async def test_multiple_sources_use_router(self, mock_memory, mock_manager): + async def test_multiple_sources_use_router( + self, mock_memory, mock_manager, mock_nucliadb_driver + ): + mock_manager.drivers.update( + { + "first-kb": mock_nucliadb_driver, + "second-kb": mock_nucliadb_driver, + } + ) expected = [MagicMock()] with patch( "hyperforge_nucliadb_agentic.agent.choose_source", @@ -63,6 +103,29 @@ async def test_multiple_sources_use_router(self, mock_memory, mock_manager): assert sources == expected choose_source.assert_awaited_once() + assert choose_source.await_args.kwargs["choose_source_model"] is None + assert choose_source.await_args.kwargs["kbid"] == "test-kbid" + + async def test_multiple_kbs_require_explicit_routing_model( + self, mock_memory, mock_manager, mock_nucliadb_driver + ): + other_driver = MagicMock() + other_driver.config.kbid = "other-kb" + mock_manager.drivers.update( + { + "first": mock_nucliadb_driver, + "second": other_driver, + } + ) + with pytest.raises(ValueError, match="explicit routing model"): + await choose_sources( + mock_memory, + mock_manager, + ["first", "second"], + "question", + ident="agent", + step_title="Choose sources", + ) class TestNucliaDBAgentInit: @@ -544,6 +607,55 @@ async def test_empty_find_returns_empty_list( class TestPrepareAskRequest: + @pytest.mark.parametrize("public_model", [None, "custom:requested-model"]) + async def test_inherits_kb_default_or_preserves_public_model( + self, mock_nucliadb_driver, public_model + ): + agent = NucliaDBAgent( + config=NucliaDBAgentConfig(sources=["kb-source-1"]), agent_id="rag-agent" + ) + public_request = ( + AskRequest(query="question", generative_model=public_model).model_dump_json( + exclude_unset=True + ) + if public_model is not None + else None + ) + request = await agent.prepare_ask_request( + mock_nucliadb_driver, "question", public_request, [] + ) + assert request.generative_model == public_model + if public_request is None: + assert "generative_model" not in request.model_fields_set + + async def test_source_search_config_supplies_model(self, mock_nucliadb_driver): + from nucliadb_models.configuration import AskConfig, AskSearchConfiguration + + agent = NucliaDBAgent( + config=NucliaDBAgentConfig( + sources=["kb-source-1"], search_config="legal-rag" + ), + agent_id="rag-agent", + ) + with patch( + "hyperforge_nucliadb_agentic.ask.search.rpc.get_search_configuration", + new=AsyncMock( + return_value=AskSearchConfiguration( + kind="ask", config=AskConfig(generative_model="custom:legal-model") + ) + ), + ): + request = await agent.prepare_ask_request( + mock_nucliadb_driver, "question", None, [] + ) + assert request.generative_model == "custom:legal-model" + + async def test_missing_kb_scope_fails(self, mock_nucliadb_driver): + mock_nucliadb_driver.config.kbid = "" + agent = NucliaDBAgent(config=NucliaDBAgentConfig(), agent_id="rag-agent") + with pytest.raises(ValueError, match="KB ID is required"): + await agent.prepare_ask_request(mock_nucliadb_driver, "question", None, []) + async def test_uses_agent_defaults(self, nucliadb_agent, mock_nucliadb_driver): request = await nucliadb_agent.prepare_ask_request( mock_nucliadb_driver, @@ -688,6 +800,32 @@ async def apply_config(reader_sdk, kbid, request): internal_request = mock_ask.call_args.kwargs["ask_request"] assert internal_request.top_k == 3 assert internal_request.filter_expression == source_filter + assert internal_request.generative_model is None + assert mock_ask.call_args.kwargs["kbid"] == "test-kbid" + mock_manager.require_kb_model_scope.assert_called_once_with("test-kbid") + + async def test_unscoped_backend_fails_before_ask( + self, mock_memory, mock_manager, mock_source + ): + agent = NucliaDBAgent( + config=NucliaDBAgentConfig(sources=["kb-source-1"]), agent_id="rag-agent" + ) + mock_manager.require_kb_model_scope.side_effect = ValueError( + "Inheriting a KB model requires a KB-scoped prediction backend" + ) + with ( + patch( + "hyperforge_nucliadb_agentic.agent.ask", new_callable=AsyncMock + ) as mock_ask, + pytest.raises(ValueError, match="KB-scoped prediction backend"), + ): + await agent.inner_rag( + source_obj=mock_source, + manager=mock_manager, + memory=mock_memory, + question="question", + ) + mock_ask.assert_not_awaited() async def test_uses_endpoint_ask_request_with_smart_agent_query( self, nucliadb_agent, mock_memory, mock_manager, mock_source @@ -741,6 +879,49 @@ async def test_uses_endpoint_ask_request_with_smart_agent_query( # --------------------------------------------------------------------------- +class TestHelperModels: + @pytest.mark.parametrize( + "method", ["inner_facets_search", "inner_facets", "inner_catalog_search"] + ) + @pytest.mark.parametrize("model", [None, "custom:kb-model"]) + async def test_helper_calls_use_source_model_and_kb( + self, + method, + model, + mock_memory, + mock_manager, + mock_source, + mock_nucliadb_driver, + ): + agent = NucliaDBAgent( + config=NucliaDBAgentConfig(sources=["kb-source-1"], generative_model=model), + agent_id="source-agent", + ) + mock_nucliadb_driver.labels.return_value = {"topic": ["legal"]} + mock_nucliadb_driver.field_labels.return_value = ({"legal": 1}, 1) + mock_nucliadb_driver.catalog_search_raw.return_value = MagicMock( + resources={}, fulltext=MagicMock(facets={"topic": {"legal": 1}}, total=0) + ) + mock_manager.execute_json.return_value = ( + {"label_sets": ["topic"], "labels": ["topic/legal"], "filters": "[]"}, + 10, + 5, + ) + + await getattr(agent, method)( + mock_memory, mock_manager, "Which documents?", mock_source, 0 + ) + + mock_manager.execute_json.assert_awaited_once() + mock_manager.execute.assert_awaited_once() + for call in [ + mock_manager.execute_json.await_args, + mock_manager.execute.await_args, + ]: + assert call.kwargs["model"] == agent.config.generative_model + assert call.kwargs["kbid"] == "test-kbid" + + class TestAllImagesByTitle: async def test_returns_empty_contexts_when_no_resources( self, nucliadb_agent, mock_memory, mock_manager, mock_nucliadb_driver diff --git a/agents/nucliadb/tests/test_config.py b/agents/nucliadb/tests/test_config.py index 6182587..ebf9d30 100644 --- a/agents/nucliadb/tests/test_config.py +++ b/agents/nucliadb/tests/test_config.py @@ -16,7 +16,7 @@ def test_sources_default_is_empty(self): def test_default_generative_model(self): cfg = NucliaDBAgentConfig() - assert cfg.generative_model.model_id == "chatgpt-azure-5.6-luna" + assert cfg.generative_model is None def test_generate_inner_answer_defaults_to_true(self): cfg = NucliaDBAgentConfig() @@ -91,6 +91,11 @@ def test_search_images_is_published(self): class TestNucliaDBAgentConfigSerialization: + def test_inherited_model_survives_round_trip(self): + cfg = NucliaDBAgentConfig(sources=["kb-a"]) + restored = NucliaDBAgentConfig.model_validate_json(cfg.model_dump_json()) + assert restored.generative_model is None + def test_round_trip_json(self): cfg = NucliaDBAgentConfig( sources=["kb-a"], diff --git a/tests/db/test_transform.py b/tests/db/test_transform.py index eff0318..423bafe 100644 --- a/tests/db/test_transform.py +++ b/tests/db/test_transform.py @@ -473,6 +473,7 @@ async def test_transform_wires_sources_and_internal_nucliadb_driver( assert isinstance(perplexity_agent, PerplexityAgentConfig) assert kb_agent.sources[0] in drivers assert kb_agent.search_config == "legal-rag" + assert kb_agent.generative_model is None assert mcp_agent.source in drivers kb_driver = drivers[kb_agent.sources[0]] mcp_driver = drivers[mcp_agent.source]