hyperforge-nucliadb-agentic 1.0.0.post96__tar.gz → 1.0.0.post108__tar.gz

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyperforge_nucliadb_agentic
-Version: 1.0.0.post96
+Version: 1.0.0.post108
 Summary: NucliaDB Hyperforge agent
 Author: Nuclia
 Author-email: Nuclia <nucliadb@nuclia.com>

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "hyperforge_nucliadb_agentic"
-version = "1.0.0.post96"
+version = "1.0.0.post108"
 license = "Apache-2.0"
 description = "NucliaDB Hyperforge agent"
 readme = "README.md"

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/pyproject.toml.orig

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "hyperforge_nucliadb_agentic"
-version = "1.0.0.post96"
+version = "1.0.0.post108"
 license = "Apache-2.0"
 description = "NucliaDB Hyperforge agent"
 authors = [{ name = "Nuclia", email = "nucliadb@nuclia.com" }]

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/src/hyperforge_nucliadb_agentic/agent.py

@@ -55,9 +55,31 @@ from hyperforge_nucliadb_agentic.ask.model import (
     RagStrategies,
     SyncAskResponse,
 )
+from hyperforge_nucliadb_agentic.ask.search import rpc
 from hyperforge_nucliadb_agentic.ask.search.ask import ask
 from hyperforge_nucliadb_agentic.config import NucliaDBAgentConfig
 
+
+async def choose_sources(
+    memory: QuestionMemory,
+    manager: Manager,
+    sources: List[str],
+    question: str,
+    ident: str,
+    step_title: str,
+) -> list[Source]:
+    if len(sources) == 1:
+        return [Source.model_construct(id=sources[0])]
+    return await choose_source(
+        memory,
+        manager,
+        sources,
+        question,
+        ident=ident,
+        step_title=step_title,
+    )
+
+
 # Example filter expressions for catalog search
 EXAMPLE_FILTER_EXP1 = [
     {
@@ -746,7 +768,7 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
                 or_filters = [f"/l/{labelset}/{lv}" for lv in label_values]
 
         # Choose sources based on the question/s and the in
-        chosen_sources = await choose_source(
+        chosen_sources = await choose_sources(
             memory,
             manager,
             sources,
@@ -783,7 +805,7 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
     ) -> List[tuple[str, str]]:
         sources = self.config.sources
         # Choose sources based on the question/s and the in
-        chosen_sources = await choose_source(
+        chosen_sources = await choose_sources(
             memory,
             manager,
             sources,
@@ -949,6 +971,57 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
         else:
             return combine_catalog_filter_expressions(to_combine, operator="and")
 
+    async def prepare_ask_request(
+        self,
+        nucliadb_driver: NucliaDBDriver,
+        question: str,
+        ask_request_json: str | None,
+        rag_strategies: list[RagStrategies],
+    ) -> AskRequest:
+        """Apply agent defaults, search config, public ASK overrides, and runtime query."""
+        # Load explicit options from the public ASK request.
+        ask_request: AskRequest | None = None
+        if ask_request_json:
+            try:
+                ask_request = AskRequest.model_validate_json(ask_request_json)
+            except ValidationError as e:
+                logger.error(
+                    f"Failed to validate AskRequest received as memory argument: {e}"
+                )
+
+        # The SmartAgent query always wins.
+        if ask_request is None:
+            ask_request = AskRequest(query=question)
+        else:
+            ask_request = ask_request.model_copy(update={"query": question})
+
+        # Search config fills fields not explicitly set by the public ASK.
+        if self.config.search_config is not None:
+            ask_request = ask_request.model_copy(
+                update={"search_configuration": self.config.search_config}
+            )
+            ask_request = await rpc.apply_ask_search_configuration(
+                nucliadb_driver.driver,
+                nucliadb_driver.config.kbid,
+                ask_request,
+            )
+
+        # Agent defaults fill any remaining fields.
+        fallback_values = {
+            "show": [ResourceProperties.BASIC, ResourceProperties.ORIGIN],
+            "citations": CitationsType.LLM_FOOTNOTES,
+            "generative_model": self.config.generative_model,
+            "rag_strategies": rag_strategies,
+            "generate_answer": self.config.generate_inner_answer,
+        }
+        return ask_request.model_copy(
+            update={
+                field: value
+                for field, value in fallback_values.items()
+                if field not in ask_request.model_fields_set
+            }
+        )
+
     async def inner_rag(
         self,
         source_obj: Source,
@@ -992,21 +1065,12 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
                 MetadataExtensionStrategy(types=["classification_labels", "origin"]),  # type: ignore
             ]
 
-        ask_request_json = memory.arguments.get("ask_request")
-        if ask_request_json:
-            # Preserve the request options from the public ask endpoint while
-            # using the query selected by the SmartAgent for this retrieval.
-            try:
-                ask_request = AskRequest.model_validate_json(
-                    ask_request_json
-                ).model_copy(update={"query": question})
-            except ValidationError as e:
-                logger.error(
-                    f"Failed to validate AskRequest received as memory argument: {e}"
-                )
-                ask_request = None
-        else:
-            ask_request = None
+        ask_request = await self.prepare_ask_request(
+            nucliadb_driver,
+            question,
+            memory.arguments.get("ask_request"),
+            rag_strategies,
+        )
 
         filter_expression = await self.build_filter_expression(
             nucliadb_driver,
@@ -1015,22 +1079,9 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
             and_filters=and_filters,
             or_filters=or_filters,
             resource_filters=resource_filters,
-            filter_expression=ask_request.filter_expression
-            if ask_request is not None
-            else None,
+            filter_expression=ask_request.filter_expression,
         )
-        if ask_request is None:
-            ask_request = AskRequest(
-                query=question,
-                show=[ResourceProperties.BASIC, ResourceProperties.ORIGIN],
-                citations=CitationsType.LLM_FOOTNOTES,
-                generative_model=self.config.generative_model,
-                filter_expression=filter_expression,
-                rag_strategies=rag_strategies,
-                generate_answer=self.config.generate_inner_answer,
-            )
-        else:
-            ask_request.filter_expression = filter_expression
+        ask_request.filter_expression = filter_expression
 
         await memory.add_step(
             step_module="nucliadb_agent",
@@ -1174,7 +1225,7 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
     ) -> List[Context]:
         sources = self.config.sources
         # Perform catalog faceted search
-        chosen_sources = await choose_source(
+        chosen_sources = await choose_sources(
             memory,
             manager,
             sources,
@@ -1333,7 +1384,7 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
     ) -> List[Context]:
         sources = self.config.sources
         # Perform catalog faceted search
-        chosen_sources = await choose_source(
+        chosen_sources = await choose_sources(
             memory,
             manager,
             sources,
@@ -1453,7 +1504,7 @@ class NucliaDBAgent(ContextAgent, Agent[NucliaDBAgentConfig]):
     ) -> List[Context]:
         sources = self.config.sources
         # Perform catalog search
-        chosen_sources = await choose_source(
+        chosen_sources = await choose_sources(
             memory,
             manager,
             sources,

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/src/hyperforge_nucliadb_agentic/ask/audit.py

@@ -187,7 +187,6 @@ class StreamAuditStorage:
             nats_subject=cast(str, usage_settings.usage_jetstream_subject),
         )
         await self.kb_usage_utility.initialize()
-
         self.initialized = True
 
     async def finalize(self):
@@ -214,8 +213,14 @@ class StreamAuditStorage:
             for event in step.external_usage or []
         ]
         if not predicts or self.kb_usage_utility is None:
+            logger.warning(
+                "Skipping KB usage report account=%s kb=%s has_predicts=%s utility_available=%s",
+                account_id,
+                kbid,
+                bool(predicts),
+                self.kb_usage_utility is not None,
+            )
             return
-
         self.kb_usage_utility.send_kb_usage(
             service=Service.RAO,
             account_id=account_id,

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/src/hyperforge_nucliadb_agentic/ask/search/rpc.py

@@ -2,7 +2,7 @@ import base64
 from typing import Literal
 
 from nucliadb_models.augment import AugmentRequest, AugmentResponse
-from nucliadb_models.configuration import SearchConfiguration
+from nucliadb_models.configuration import AskConfig, SearchConfiguration
 from nucliadb_models.graph.requests import GraphNodesSearchRequest, GraphSearchRequest
 from nucliadb_models.graph.responses import (
     GraphNodesSearchResponse,
@@ -25,7 +25,7 @@ from hyperforge_nucliadb_agentic.ask.exceptions import (
     KnowledgeBoxNotFound,
     NucliaDBError,
 )
-from hyperforge_nucliadb_agentic.ask.model import Image
+from hyperforge_nucliadb_agentic.ask.model import AskRequest, Image
 from hyperforge_nucliadb_agentic.ask.settings import settings
 from hyperforge_nucliadb_agentic.ask.utils.ids import FieldId
 
@@ -60,6 +60,36 @@ async def get_search_configuration(
     return config
 
 
+class SearchConfigurationNotFound(Exception):
+    pass
+
+
+class InvalidAskSearchConfiguration(Exception):
+    pass
+
+
+async def apply_ask_search_configuration(
+    reader_sdk: NucliaDBAsync,
+    kbid: str,
+    ask_request: AskRequest,
+) -> AskRequest:
+    if ask_request.search_configuration is None:
+        return ask_request
+
+    search_config = await get_search_configuration(
+        reader_sdk, kbid, name=ask_request.search_configuration
+    )
+    if search_config is None:
+        raise SearchConfigurationNotFound
+    if not isinstance(search_config.config, AskConfig):
+        raise InvalidAskSearchConfiguration
+
+    return AskRequest.model_validate(
+        search_config.config.model_dump(exclude_unset=True)
+        | ask_request.model_dump(exclude_unset=True)
+    )
+
+
 async def find(
     search_sdk: NucliaDBAsync,
     kbid: str,

{hyperforge_nucliadb_agentic-1.0.0.post96 → hyperforge_nucliadb_agentic-1.0.0.post108}/src/hyperforge_nucliadb_agentic/config.py

@@ -9,6 +9,10 @@ from pydantic.config import ConfigDict
 class NucliaDBAgentConfig(ContextAgentConfig):
     model_config = ConfigDict(title="Knowledge Box Agent")
     module: Literal["nucliadb_agent"] = "nucliadb_agent"
+    search_config: Optional[str] = Field(
+        default=None,
+        description="ID of an ask search configuration applied to this agent",
+    )
     sources: List[str] = Field(
         default_factory=list,
         json_schema_extra={