nucliadb 6.2.0.post2675__py3-none-any.whl → 6.2.1__py3-none-any.whl

nucliadb/search/search/hydrator.py

@@ -66,6 +66,9 @@ class TextBlockHydrationOptions(BaseModel):
     # list of exact matches to highlight
     ematches: Optional[list[str]] = None
 
+    # If true, only hydrate the text block if its text is not already populated
+    only_hydrate_empty: bool = False
+
 
 @hydrator_observer.wrap({"type": "resource_text"})
 async def hydrate_resource_text(
@@ -161,6 +164,8 @@ async def hydrate_text_block(
     `text_block` object.
 
     """
+    if options.only_hydrate_empty and text_block.text:
+        return text_block
     async with AsyncExitStack() as stack:
         if concurrency_control is not None:
             await stack.enter_async_context(concurrency_control)
@@ -188,4 +193,5 @@ def text_block_to_find_paragraph(text_block: TextBlockMatch) -> FindParagraph:
         reference=text_block.representation_file,
         page_with_visual=text_block.page_with_visual,
         position=text_block.position,
+        relevant_relations=text_block.relevant_relations,
     )

nucliadb/search/search/merge.py

@@ -23,6 +23,8 @@ import math
 from typing import Any, Optional, Set, Union
 
 from nucliadb.common.ids import FieldId, ParagraphId
+from nucliadb.common.models_utils import from_proto
+from nucliadb.common.models_utils.from_proto import RelationTypePbMap
 from nucliadb.search.search import cache
 from nucliadb.search.search.cut import cut_page
 from nucliadb.search.search.fetch import (
@@ -33,11 +35,11 @@ from nucliadb.search.search.fetch import (
 )
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.labels import translate_system_to_alias_label
-from nucliadb_models.metadata import RelationTypePbMap
 from nucliadb_models.resource import ExtractedDataTypeName
 from nucliadb_models.search import (
     DirectionalRelation,
     EntitySubgraph,
+    EntityType,
     KnowledgeboxSearchResults,
     KnowledgeboxSuggestResults,
     MinScore,
@@ -46,7 +48,6 @@ from nucliadb_models.search import (
     RelatedEntities,
     RelatedEntity,
     RelationDirection,
-    RelationNodeTypeMap,
     Relations,
     ResourceProperties,
     ResourceResult,
@@ -71,6 +72,7 @@ from nucliadb_protos.nodereader_pb2 import (
     SuggestResponse,
     VectorSearchResponse,
 )
+from nucliadb_protos.utils_pb2 import RelationNode
 
 from .metrics import merge_observer
 from .paragraphs import get_paragraph_text, get_text_sentence
@@ -81,6 +83,15 @@ TitleScore = str
 SortValue = Union[Bm25Score, TimestampScore, TitleScore]
 
 
+def relation_node_type_to_entity_type(node_type: RelationNode.NodeType.ValueType) -> EntityType:
+    return {
+        RelationNode.NodeType.ENTITY: EntityType.ENTITY,
+        RelationNode.NodeType.LABEL: EntityType.LABEL,
+        RelationNode.NodeType.RESOURCE: EntityType.RESOURCE,
+        RelationNode.NodeType.USER: EntityType.USER,
+    }[node_type]
+
+
 def sort_results_by_score(results: Union[list[ParagraphResult], list[DocumentResult]]):
     results.sort(key=lambda x: (x.score.bm25, x.score.booster), reverse=True)
 
@@ -432,15 +443,38 @@ async def merge_paragraph_results(
 async def merge_relations_results(
     relations_responses: list[RelationSearchResponse],
     query: EntitiesSubgraphRequest,
+    only_with_metadata: bool = False,
+    only_agentic: bool = False,
 ) -> Relations:
     loop = asyncio.get_event_loop()
-    return await loop.run_in_executor(None, _merge_relations_results, relations_responses, query)
+    return await loop.run_in_executor(
+        None,
+        _merge_relations_results,
+        relations_responses,
+        query,
+        only_with_metadata,
+        only_agentic,
+    )
 
 
 def _merge_relations_results(
     relations_responses: list[RelationSearchResponse],
     query: EntitiesSubgraphRequest,
+    only_with_metadata: bool,
+    only_agentic: bool,
 ) -> Relations:
+    """
+    Merge relation search responses into a single Relations object while applying filters.
+
+    Args:
+        relations_responses: List of relation search responses
+        query: EntitiesSubgraphRequest object
+        only_with_metadata: If True, only include relations with metadata. This metadata includes paragraph_id and entity positions among other things.
+        only_agentic: If True, only include relations extracted by a Graph Extraction Agent.
+
+    Returns:
+        Relations
+    """
     relations = Relations(entities={})
 
     for entry_point in query.entry_points:
@@ -452,27 +486,37 @@ def _merge_relations_results(
             destination = relation.to
             relation_type = RelationTypePbMap[relation.relation]
             relation_label = relation.relation_label
-
-            if origin.value in relations.entities:
-                relations.entities[origin.value].related_to.append(
-                    DirectionalRelation(
-                        entity=destination.value,
-                        entity_type=RelationNodeTypeMap[destination.ntype],
-                        relation=relation_type,
-                        relation_label=relation_label,
-                        direction=RelationDirection.OUT,
+            metadata = relation.metadata if relation.HasField("metadata") else None
+            # If only_with_metadata is True, we check that metadata for the relation is not None
+            # If only_agentic is True, we check that metadata for the relation is not None and that it has a data_augmentation_task_id
+            # TODO: This is suboptimal, we should be able to filter this in the query to the index,
+            if (not only_with_metadata or metadata) and (
+                not only_agentic or (metadata and metadata.data_augmentation_task_id)
+            ):
+                if origin.value in relations.entities:
+                    relations.entities[origin.value].related_to.append(
+                        DirectionalRelation(
+                            entity=destination.value,
+                            entity_type=relation_node_type_to_entity_type(destination.ntype),
+                            entity_subtype=destination.subtype,
+                            relation=relation_type,
+                            relation_label=relation_label,
+                            direction=RelationDirection.OUT,
+                            metadata=from_proto.relation_metadata(metadata) if metadata else None,
+                        )
                     )
-                )
-            elif destination.value in relations.entities:
-                relations.entities[destination.value].related_to.append(
-                    DirectionalRelation(
-                        entity=origin.value,
-                        entity_type=RelationNodeTypeMap[origin.ntype],
-                        relation=relation_type,
-                        relation_label=relation_label,
-                        direction=RelationDirection.IN,
+                elif destination.value in relations.entities:
+                    relations.entities[destination.value].related_to.append(
+                        DirectionalRelation(
+                            entity=origin.value,
+                            entity_type=relation_node_type_to_entity_type(origin.ntype),
+                            entity_subtype=origin.subtype,
+                            relation=relation_type,
+                            relation_label=relation_label,
+                            direction=RelationDirection.IN,
+                            metadata=from_proto.relation_metadata(metadata) if metadata else None,
+                        )
                     )
-                )
 
     return relations
 
@@ -571,11 +615,22 @@ async def merge_suggest_entities_results(
     return RelatedEntities(entities=list(unique_entities), total=len(unique_entities))
 
 
+def merge_relation_prefix_results(
+    responses: list[SearchResponse],
+) -> RelatedEntities:
+    unique_entities: Set[RelatedEntity] = set()
+    for response in responses:
+        response_entities = (
+            RelatedEntity(family=e.subtype, value=e.value) for e in response.relation.prefix.nodes
+        )
+        unique_entities.update(response_entities)
+
+    return RelatedEntities(entities=list(unique_entities), total=len(unique_entities))
+
+
 async def merge_suggest_results(
     suggest_responses: list[SuggestResponse],
     kbid: str,
-    show: list[ResourceProperties],
-    field_type_filter: list[FieldTypeName],
     highlight: bool = False,
 ) -> KnowledgeboxSuggestResults:
     api_results = KnowledgeboxSuggestResults()