Return NER family in entity suggestions (#1729)

Return NER family in entity suggestions

.gitignore

@@ -23,6 +23,7 @@ target
 .egg-info
 *.so
 venv/
+.venv/
 nucliadb/.mypy_cache
 lib
 share

nucliadb/nucliadb/search/search/merge.py

@@ -19,7 +19,7 @@
 #
 import datetime
 import math
-from typing import Any, Optional, Union
+from typing import Any, Optional, Set, Union
 
 from nucliadb_protos.nodereader_pb2 import (
     DocumentResult,
@@ -53,6 +53,7 @@
     Paragraph,
     Paragraphs,
     RelatedEntities,
+    RelatedEntity,
     RelationDirection,
     RelationNodeTypeMap,
     Relations,
@@ -594,13 +595,15 @@ async def merge_paragraphs_results(
 async def merge_suggest_entities_results(
     suggest_responses: list[SuggestResponse],
 ) -> RelatedEntities:
-    merge = RelatedEntities(entities=[], total=0)
-
+    unique_entities: Set[RelatedEntity] = set()
     for response in suggest_responses:
-        merge.entities.extend(response.entities.entities)
-        merge.total += response.entities.total
+        response_entities = (
+            RelatedEntity(family=e.subtype, value=e.value)
+            for e in response.entity_results.nodes
+        )
+        unique_entities.update(response_entities)
 
-    return merge
+    return RelatedEntities(entities=list(unique_entities), total=len(unique_entities))
 
 
 async def merge_suggest_results(

nucliadb/nucliadb/tests/integration/test_entities.py

@@ -431,7 +431,7 @@ async def test_entities_indexing(
     assert resp.status_code == 200
     body = resp.json()
 
-    entities = set(body["entities"]["entities"])
+    entities = set((e["value"] for e in body["entities"]["entities"]))
     # BUG? why is "domestic cat" not appearing in the results?
     assert entities == {"dog", "dolphin"}
     # assert entities == {"dog", "domestic cat", "dolphin"}

nucliadb/nucliadb/tests/integration/test_suggest.py

@@ -212,16 +212,19 @@ async def test_suggest_related_entities(
     )
     assert resp.status_code == 201
 
+    def assert_expected_entities(body, expected):
+        assert set((e["value"] for e in body["entities"]["entities"])) == expected
+
     # Test simple suggestions
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=Ann")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Anna", "Anthony"}
+    assert_expected_entities(body, {"Anna", "Anthony"})
 
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=joh")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"John"}
+    assert_expected_entities(body, {"John"})
 
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=xxxxx")
     assert resp.status_code == 200
@@ -233,33 +236,33 @@ async def test_suggest_related_entities(
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=bar")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Barcelona", "Bárcenas"}
+    assert_expected_entities(body, {"Barcelona", "Bárcenas"})
 
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=Bar")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Barcelona", "Bárcenas"}
+    assert_expected_entities(body, {"Barcelona", "Bárcenas"})
 
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=BAR")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Barcelona", "Bárcenas"}
+    assert_expected_entities(body, {"Barcelona", "Bárcenas"})
 
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=BÄR")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Barcelona", "Bárcenas"}
+    assert_expected_entities(body, {"Barcelona", "Bárcenas"})
 
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=BáR")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Barcelona", "Bárcenas"}
+    assert_expected_entities(body, {"Barcelona", "Bárcenas"})
 
     # Test multiple word suggest and ordering
     resp = await nucliadb_reader.get(f"/kb/{knowledgebox}/suggest?query=Solomon+Is")
     assert resp.status_code == 200
     body = resp.json()
-    assert set(body["entities"]["entities"]) == {"Solomon Islands", "Israel"}
+    assert_expected_entities(body, {"Solomon Islands", "Israel"})
 
 
 @pytest.mark.asyncio
@@ -350,8 +353,8 @@ def assert_expected_paragraphs(response):
 
     def assert_expected_entities(response):
         expected = {"Anna", "Anthony"}
-        assert response["entities"]["total"] == len(expected)
-        assert set(response["entities"]["entities"]) == expected
+        assert len(response["entities"]) == len(expected)
+        assert set((e["value"] for e in response["entities"]["entities"])) == expected
 
     resp = await nucliadb_reader.get(
         f"/kb/{knowledgebox}/suggest",

nucliadb_models/nucliadb_models/search.py

@@ -228,9 +228,14 @@ class Relations(BaseModel):
     # graph: List[RelationPath]
 
 
+class RelatedEntity(BaseModel, frozen=True):
+    family: str
+    value: str
+
+
 class RelatedEntities(BaseModel):
     total: int = 0
-    entities: List[str] = []
+    entities: List[RelatedEntity] = []
 
 
 class ResourceSearchResults(JsonBaseModel):

nucliadb_node/src/shards/shard_reader.rs

@@ -29,7 +29,7 @@ use nucliadb_core::protos::shard_created::{
 };
 use nucliadb_core::protos::{
     DocumentSearchRequest, DocumentSearchResponse, EdgeList, GetShardRequest,
-    ParagraphSearchRequest, ParagraphSearchResponse, RelatedEntities, RelationPrefixSearchRequest,
+    ParagraphSearchRequest, ParagraphSearchResponse, RelationPrefixSearchRequest,
     RelationSearchRequest, RelationSearchResponse, SearchRequest, SearchResponse, Shard, ShardFile,
     ShardFileChunk, ShardFileList, StreamRequest, SuggestFeatures, SuggestRequest, SuggestResponse,
     TypeList, VectorSearchRequest, VectorSearchResponse,
@@ -38,8 +38,9 @@ use nucliadb_core::query_planner::QueryPlan;
 use nucliadb_core::thread::*;
 use nucliadb_core::tracing::{self, *};
 use nucliadb_procs::measure;
-use nucliadb_protos::nodereader::RelationNodeFilter;
+use nucliadb_protos::nodereader::{RelationNodeFilter, RelationPrefixSearchResponse};
 use nucliadb_protos::utils::relation_node::NodeType;
+use nucliadb_relations2::reader::HashedRelationNode;
 
 use crate::disk_structure::*;
 use crate::shards::metadata::ShardMetadata;
@@ -362,13 +363,12 @@ impl ShardReader {
                     .into_iter()
                     .flatten() // unwrap errors and continue with successful results
                     .flat_map(|response| response.prefix)
-                    .flat_map(|prefix_response| prefix_response.nodes.into_iter())
-                    .map(|node| node.value);
+                    .flat_map(|prefix_response| prefix_response.nodes.into_iter());
 
                 // remove duplicate entities
-                let mut seen = HashSet::new();
+                let mut seen: HashSet<HashedRelationNode> = HashSet::new();
                 let mut ent_result = entities.collect::<Vec<_>>();
-                ent_result.retain(|e| seen.insert(e.clone()));
+                ent_result.retain(|e| seen.insert(e.clone().into()));
 
                 ent_result
             };
@@ -403,10 +403,7 @@ impl ShardReader {
         };
 
         if let Some(entities) = entities {
-            response.entities = Some(RelatedEntities {
-                total: entities.len() as u32,
-                entities,
-            })
+            response.entity_results = Some(RelationPrefixSearchResponse { nodes: entities });
         }
 
         Ok(response)

nucliadb_node/tests/test_suggest.rs

@@ -203,10 +203,10 @@ async fn test_suggest_entities(
 
     // multiple words and result ordering
     let response = suggest_entities(&mut reader, &shard.id, "Solomon Isa").await;
-    assert!(response.entities.is_some());
-    assert_eq!(response.entities.as_ref().unwrap().total, 2);
-    assert!(response.entities.as_ref().unwrap().entities[0] == *"Solomon Islands");
-    assert!(response.entities.as_ref().unwrap().entities[1] == *"Israel");
+    assert!(response.entity_results.is_some());
+    assert_eq!(response.entity_results.as_ref().unwrap().nodes.len(), 2);
+    assert!(response.entity_results.as_ref().unwrap().nodes[0].value == *"Solomon Islands");
+    assert!(response.entity_results.as_ref().unwrap().nodes[1].value == *"Israel");
 
     // Does not find resources by UUID prefix
     let pap_uuid = &shard.resources["pap"];
@@ -238,7 +238,7 @@ async fn test_suggest_features(
     let mut reader = node.reader_client();
 
     let response = suggest_paragraphs(&mut reader, &shard.id, "ann").await;
-    assert!(response.entities.is_none());
+    assert!(response.entity_results.is_none());
     expect_paragraphs(
         &response,
         &[(&shard.resources["little prince"], "/a/summary")],
@@ -358,20 +358,21 @@ async fn suggest_entities(
 }
 
 fn expect_entities(response: &SuggestResponse, expected: &[&str]) {
-    assert!(response.entities.is_some());
+    assert!(response.entity_results.is_some());
     assert_eq!(
-        response.entities.as_ref().unwrap().total as usize,
+        response.entity_results.as_ref().unwrap().nodes.len(),
         expected.len(),
         "Response entities don't match expected ones: {:?} != {:?}",
-        response.entities,
+        response.entity_results,
         expected,
     );
     for entity in expected {
         assert!(response
-            .entities
+            .entity_results
             .as_ref()
             .unwrap()
-            .entities
-            .contains(&entity.to_string()));
+            .nodes
+            .iter()
+            .any(|e| &e.value == entity));
     }
 }

nucliadb_protos/nodereader.proto

@@ -368,7 +368,7 @@ message SuggestResponse {
     repeated string ematches = 4;
 
     // Entities related with the query
-    RelatedEntities entities = 5;
+    RelationPrefixSearchResponse entity_results = 6;
 }
 
 message SearchResponse {