nucliadb 6.2.1.post3063__py3-none-any.whl → 6.2.1.post3071__py3-none-any.whl

nucliadb/common/cluster/manager.py

@@ -104,6 +104,14 @@ class KBShardManager:
             errors.capture_exception(exc)
             raise NodeError("Node unavailable for operation") from exc
 
+        for result in results:
+            if isinstance(result, Exception):
+                errors.capture_exception(result)
+                raise NodeError(
+                    f"Error while applying {aw.__name__} for all shards. Other similar errors may have been shadowed.\n"
+                    f"{type(result).__name__}: {result}"
+                ) from result
+
         return results
 
     # TODO: move to data manager

nucliadb/ingest/orm/entities.py

@@ -26,7 +26,6 @@ from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.cluster.exceptions import (
     AlreadyExists,
     EntitiesGroupNotFound,
-    NodeError,
 )
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.datamanagers.entities import (
@@ -54,7 +53,6 @@ from nucliadb_protos.nodereader_pb2 import (
 )
 from nucliadb_protos.utils_pb2 import RelationNode
 from nucliadb_protos.writer_pb2 import GetEntitiesResponse
-from nucliadb_telemetry import errors
 
 from .exceptions import EntityManagementException
 
@@ -226,10 +224,6 @@ class EntitiesManager:
             settings.relation_search_timeout,
             use_read_replica_nodes=self.use_read_replica_nodes,
         )
-        for result in results:
-            if isinstance(result, Exception):
-                errors.capture_exception(result)
-                raise NodeError("Error while querying relation index")
 
         entities = {}
         for result in results:
@@ -305,6 +299,7 @@ class EntitiesManager:
         shard_manager = get_shard_manager()
 
         async def query_indexed_entities_group_names(node: AbstractIndexNode, shard_id: str) -> set[str]:
+            """Search all relation types"""
             request = SearchRequest(
                 shard=shard_id,
                 result_per_page=0,
@@ -316,10 +311,11 @@ class EntitiesManager:
             response: SearchResponse = await query_shard(node, shard_id, request)
             try:
                 facetresults = response.document.facets["/e"].facetresults
-                return {facet.tag.split("/")[-1] for facet in facetresults}
             except KeyError:
                 # No entities found
                 return set()
+            else:
+                return {facet.tag.split("/")[-1] for facet in facetresults}
 
         results = await shard_manager.apply_for_all_shards(
             self.kbid,
@@ -327,10 +323,6 @@ class EntitiesManager:
             settings.relation_types_timeout,
             use_read_replica_nodes=self.use_read_replica_nodes,
         )
-        for result in results:
-            if isinstance(result, Exception):
-                errors.capture_exception(result)
-                raise NodeError("Error while looking for relations types")
 
         if not results:
             return set()

nucliadb/search/search/chat/ask.py

@@ -719,9 +719,19 @@ async def retrieval_in_kb(
     prequeries = parse_prequeries(ask_request)
     graph_strategy = parse_graph_strategy(ask_request)
     with metrics.time("retrieval"):
-        prequeries_results = None
+        main_results, prequeries_results, query_parser = await get_find_results(
+            kbid=kbid,
+            query=main_query,
+            item=ask_request,
+            ndb_client=client_type,
+            user=user_id,
+            origin=origin,
+            metrics=metrics,
+            prequeries_strategy=prequeries,
+        )
+
         if graph_strategy is not None:
-            main_results, query_parser = await get_graph_results(
+            graph_results, graph_request = await get_graph_results(
                 kbid=kbid,
                 query=main_query,
                 item=ask_request,
@@ -732,18 +742,13 @@ async def retrieval_in_kb(
                 metrics=metrics,
                 shards=ask_request.shards,
             )
-        # TODO (oni): Fallback to normal retrieval if no graph results are found
-        else:
-            main_results, prequeries_results, query_parser = await get_find_results(
-                kbid=kbid,
-                query=main_query,
-                item=ask_request,
-                ndb_client=client_type,
-                user=user_id,
-                origin=origin,
-                metrics=metrics,
-                prequeries_strategy=prequeries,
-            )
+
+            if prequeries_results is None:
+                prequeries_results = []
+
+            prequery = PreQuery(id="graph", request=graph_request, weight=graph_strategy.weight)
+            prequeries_results.append((prequery, graph_results))
+
         if len(main_results.resources) == 0 and all(
             len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
         ):

nucliadb/search/search/graph_strategy.py

@@ -17,11 +17,9 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
-import asyncio
 import heapq
 import json
 from collections import defaultdict
-from datetime import datetime
 from typing import Any, Collection, Iterable, Optional, Union
 
 from nuclia_models.predict.generative_responses import (
@@ -46,9 +44,8 @@ from nucliadb.search.search.find_merge import (
     hydrate_and_rerank,
 )
 from nucliadb.search.search.hydrator import ResourceHydrationOptions, TextBlockHydrationOptions
-from nucliadb.search.search.merge import merge_suggest_results
+from nucliadb.search.search.merge import merge_relation_prefix_results
 from nucliadb.search.search.metrics import RAGMetrics
-from nucliadb.search.search.query import QueryParser
 from nucliadb.search.search.rerankers import Reranker, RerankingOptions
 from nucliadb.search.utilities import get_predict
 from nucliadb_models.common import FieldTypeName
@@ -62,11 +59,12 @@ from nucliadb_models.search import (
     ChatModel,
     DirectionalRelation,
     EntitySubgraph,
+    FindRequest,
     GraphStrategy,
     KnowledgeboxFindResults,
-    KnowledgeboxSuggestResults,
     NucliaDBClientType,
     QueryEntityDetection,
+    RelatedEntities,
     RelationDirection,
     RelationRanking,
     Relations,
@@ -308,7 +306,7 @@ async def get_graph_results(
     generative_model: Optional[str] = None,
     metrics: RAGMetrics = RAGMetrics(),
     shards: Optional[list[str]] = None,
-) -> tuple[KnowledgeboxFindResults, QueryParser]:
+) -> tuple[KnowledgeboxFindResults, FindRequest]:
     relations = Relations(entities={})
     explored_entities: set[str] = set()
     scores: dict[str, list[float]] = {}
@@ -321,23 +319,18 @@ async def get_graph_results(
             # Get the entities from the query
             with metrics.time("graph_strat_query_entities"):
                 if graph_strategy.query_entity_detection == QueryEntityDetection.SUGGEST:
-                    suggest_result = await fuzzy_search_entities(
+                    relation_result = await fuzzy_search_entities(
                         kbid=kbid,
                         query=query,
-                        range_creation_start=item.range_creation_start,
-                        range_creation_end=item.range_creation_end,
-                        range_modification_start=item.range_modification_start,
-                        range_modification_end=item.range_modification_end,
-                        target_shard_replicas=shards,
                     )
-                    if suggest_result.entities is not None:
+                    if relation_result is not None:
                         entities_to_explore = (
                             RelationNode(
                                 ntype=RelationNode.NodeType.ENTITY,
                                 value=result.value,
                                 subtype=result.family,
                             )
-                            for result in suggest_result.entities.entities
+                            for result in relation_result.entities
                         )
                 elif (
                     not entities_to_explore
@@ -361,6 +354,7 @@ async def get_graph_results(
                 for relation in subgraph.related_to
                 if relation.entity not in explored_entities
             )
+
         # Get the relations for the new entities
         with metrics.time("graph_strat_neighbor_relations"):
             try:
@@ -437,54 +431,36 @@ async def get_graph_results(
             field_type_filter=find_request.field_type_filter,
             relation_text_as_paragraphs=graph_strategy.relation_text_as_paragraphs,
         )
-    return find_results, query_parser
+    return find_results, find_request
 
 
 async def fuzzy_search_entities(
     kbid: str,
     query: str,
-    range_creation_start: Optional[datetime] = None,
-    range_creation_end: Optional[datetime] = None,
-    range_modification_start: Optional[datetime] = None,
-    range_modification_end: Optional[datetime] = None,
-    target_shard_replicas: Optional[list[str]] = None,
-) -> KnowledgeboxSuggestResults:
+) -> Optional[RelatedEntities]:
     """Fuzzy find entities in KB given a query using the same methodology as /suggest, but split by words."""
 
-    base_request = nodereader_pb2.SuggestRequest(
-        body="", features=[nodereader_pb2.SuggestFeatures.ENTITIES]
-    )
-    if range_creation_start is not None:
-        base_request.timestamps.from_created.FromDatetime(range_creation_start)
-    if range_creation_end is not None:
-        base_request.timestamps.to_created.FromDatetime(range_creation_end)
-    if range_modification_start is not None:
-        base_request.timestamps.from_modified.FromDatetime(range_modification_start)
-    if range_modification_end is not None:
-        base_request.timestamps.to_modified.FromDatetime(range_modification_end)
-
-    tasks = []
-    # XXX: Splitting by words is not ideal, in the future, modify suggest to better handle this
-    for word in query.split():
-        if len(word) < 3:
-            continue
-        request = nodereader_pb2.SuggestRequest()
-        request.CopyFrom(base_request)
-        request.body = word
-        tasks.append(
-            node_query(kbid, Method.SUGGEST, request, target_shard_replicas=target_shard_replicas)
-        )
+    request = nodereader_pb2.SearchRequest()
+    request.relation_prefix.query = query
 
+    results: list[nodereader_pb2.SearchResponse]
     try:
-        results_raw = await asyncio.gather(*tasks)
-        return await merge_suggest_results(
-            [item for r in results_raw for item in r[0]],
-            kbid=kbid,
+        (
+            results,
+            _,
+            _,
+        ) = await node_query(
+            kbid,
+            Method.SEARCH,
+            request,
+            use_read_replica_nodes=True,
+            retry_on_primary=False,
         )
+        return merge_relation_prefix_results(results)
     except Exception as e:
         capture_exception(e)
         logger.exception("Error in finding entities in query for graph strategy")
-        return KnowledgeboxSuggestResults(entities=None)
+        return None
 
 
 async def rank_relations_reranker(

nucliadb/search/search/merge.py

@@ -615,6 +615,19 @@ 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,

{nucliadb-6.2.1.post3063.dist-info → nucliadb-6.2.1.post3071.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.2.1.post3063
+Version: 6.2.1.post3071
 Home-page: https://docs.nuclia.dev/docs/management/nucliadb/intro
 Author: NucliaDB Community
 Author-email: nucliadb@nuclia.com
@@ -22,10 +22,10 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: >=3.9, <4
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.2.1.post3063
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post3063
-Requires-Dist: nucliadb-protos>=6.2.1.post3063
-Requires-Dist: nucliadb-models>=6.2.1.post3063
+Requires-Dist: nucliadb-telemetry[all]>=6.2.1.post3071
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post3071
+Requires-Dist: nucliadb-protos>=6.2.1.post3071
+Requires-Dist: nucliadb-models>=6.2.1.post3071
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn

{nucliadb-6.2.1.post3063.dist-info → nucliadb-6.2.1.post3071.dist-info}/RECORD

@@ -47,7 +47,7 @@ nucliadb/common/cluster/base.py,sha256=kklDqyvsubNX0W494ttl9f3E58lGaX6AXqAd8XX8Z
 nucliadb/common/cluster/exceptions.py,sha256=V3c_fgH00GyJ-a5CaGLhwTuhwhUNR9YAGvS5jaRuc_Y,1495
 nucliadb/common/cluster/grpc_node_dummy.py,sha256=L85wBnfab7Rev0CfsfUjPxQC6DiHPsETKrZAOLx9XHg,3510
 nucliadb/common/cluster/index_node.py,sha256=g38H1kiAliF3Y6et_CWYInpn_xPxf7THAFJ7RtgLNZo,3246
-nucliadb/common/cluster/manager.py,sha256=3UnYwVb-ZykYfLndxM7TLw7-2T_vxqoFXMu0Pzxh5-A,15327
+nucliadb/common/cluster/manager.py,sha256=cj8yNIspsmdyGUq3vbyW--r-LcxAUqA4S9tocz6MsCM,15695
 nucliadb/common/cluster/rebalance.py,sha256=jSEYsPgs_Dobv3FOaKl5arBko4s8JlWkahm8LOzgNnE,9135
 nucliadb/common/cluster/rollover.py,sha256=dx6AF9ywKP10iBNlcoJgRV40921fOPpVWaCUU54hztE,25823
 nucliadb/common/cluster/settings.py,sha256=TMoym-cZsQ2soWfLAce0moSa2XncttQyhahL43LrWTo,3384
@@ -128,7 +128,7 @@ nucliadb/ingest/fields/text.py,sha256=tFvSQJAe0W7ePpp2_WDfLiE2yglR1OTU0Zht9acvOF
 nucliadb/ingest/orm/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/ingest/orm/brain.py,sha256=UND5EsNUdd7XdjScYqRqg4r_xCx3l-My8alGw5M9CWg,28398
 nucliadb/ingest/orm/broker_message.py,sha256=ZEMueoGuuRKO4tHgzc0P0AM1Ls1TTYey_4UvRQf0BpY,6915
-nucliadb/ingest/orm/entities.py,sha256=5d6Gfo-Yz-rns_mNJeRqiGaPeWpUMgSKZnmWIGMLCKo,15537
+nucliadb/ingest/orm/entities.py,sha256=sBhg8eahsWVwO34KoAJV2YRix4Uw5GINx3srJWxRC9k,15148
 nucliadb/ingest/orm/exceptions.py,sha256=k4Esv4NtL4TrGTcsQpwrSfDhPQpiYcRbB1SpYmBX5MY,1432
 nucliadb/ingest/orm/knowledgebox.py,sha256=IGOPvBR1qXqDxE5DeiOdYCLdPgjzOVVpsASJ2zYvWwQ,23651
 nucliadb/ingest/orm/metrics.py,sha256=OkwMSPKLZcKba0ZTwtTiIxwBgaLMX5ydhGieKvi2y7E,1096
@@ -210,9 +210,9 @@ nucliadb/search/search/fetch.py,sha256=XJHIFnZmXM_8Kb37lb4lg1GYG7cZ1plT-qAIb_Qzi
 nucliadb/search/search/filters.py,sha256=1MkHlJjAQqoRCj7e5cEzK2HvBxGLE17I_omsjiklbtw,6476
 nucliadb/search/search/find.py,sha256=yQbttt85wQFc4NEaj2RNGgozP7IQx_bjAOhHke3fXY0,9890
 nucliadb/search/search/find_merge.py,sha256=_R_YpHAZv5BHh3XABQ8MRd1Ci0seclGYf26yJHJ7H0I,17178
-nucliadb/search/search/graph_strategy.py,sha256=6d-KjGDbOnaXQzEwyBpA-iQM0rkveVTiK3A3m2UJq8Q,33538
+nucliadb/search/search/graph_strategy.py,sha256=cFcu6nrOqHfobsbUu1pOwfBndrN4ppEvv3-4jV4_4bg,31977
 nucliadb/search/search/hydrator.py,sha256=-R37gCrGxkyaiHQalnTWHNG_FCx11Zucd7qA1vQCxuw,6985
-nucliadb/search/search/merge.py,sha256=g0PxUejWtYIYWG-VroArMCgwB6AOp3lZMkzoHAPYBKE,22183
+nucliadb/search/search/merge.py,sha256=i_PTBFRqC5iTTziOMEltxLIlmokIou5hjjgR4BnoLBE,22635
 nucliadb/search/search/metrics.py,sha256=81X-tahGW4n2CLvUzCPdNxNClmZqUWZjcVOGCUHoiUM,2872
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=IaNK4dAxdXs38PoIkTdgqMDuZDjeiOtcXn3LeaT-OMw,8855
@@ -224,7 +224,7 @@ nucliadb/search/search/shards.py,sha256=JSRSrHgHcF4sXyuZZoJdMfK0v_LHpoSRf1lCr5-K
 nucliadb/search/search/summarize.py,sha256=ksmYPubEQvAQgfPdZHfzB_rR19B2ci4IYZ6jLdHxZo8,4996
 nucliadb/search/search/utils.py,sha256=iF2tbBA56gRMJH1TlE2hMrqeXqjoeOPt4KgRdp2m9Ek,3313
 nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/chat/ask.py,sha256=tE1Q5V58oLMCo-T9s0N6Kko-1RWn1e4kHfbbPBsD2uU,36266
+nucliadb/search/search/chat/ask.py,sha256=K85Size6WAb-q4sCn0u1drrPnqIvqCy6YbfCxQmsEt8,36348
 nucliadb/search/search/chat/exceptions.py,sha256=Siy4GXW2L7oPhIR86H3WHBhE9lkV4A4YaAszuGGUf54,1356
 nucliadb/search/search/chat/images.py,sha256=PA8VWxT5_HUGfW1ULhKTK46UBsVyINtWWqEM1ulzX1E,3095
 nucliadb/search/search/chat/prompt.py,sha256=r2JTiRWH3YHPdeRAG5w6gD0g0fWVxdTjYIR86qAVa7k,47106
@@ -329,9 +329,9 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.2.1.post3063.dist-info/METADATA,sha256=LB2WO-kk-GSCoSk1ef5S7xI1WmB7BTeo_6WvD2RdtlM,4603
-nucliadb-6.2.1.post3063.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nucliadb-6.2.1.post3063.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.2.1.post3063.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.2.1.post3063.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb-6.2.1.post3063.dist-info/RECORD,,
+nucliadb-6.2.1.post3071.dist-info/METADATA,sha256=mbRwoP8OS4lUygXdPQgnAcU7Q3o7AunTre_KXI5E66o,4603
+nucliadb-6.2.1.post3071.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nucliadb-6.2.1.post3071.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.2.1.post3071.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.2.1.post3071.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb-6.2.1.post3071.dist-info/RECORD,,