nucliadb 6.2.1.post3059__py3-none-any.whl → 6.2.1.post3067__py3-none-any.whl

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.post3059.dist-info → nucliadb-6.2.1.post3067.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.2.1.post3059
+Version: 6.2.1.post3067
 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.post3059
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post3059
-Requires-Dist: nucliadb-protos>=6.2.1.post3059
-Requires-Dist: nucliadb-models>=6.2.1.post3059
+Requires-Dist: nucliadb-telemetry[all]>=6.2.1.post3067
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post3067
+Requires-Dist: nucliadb-protos>=6.2.1.post3067
+Requires-Dist: nucliadb-models>=6.2.1.post3067
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn

{nucliadb-6.2.1.post3059.dist-info → nucliadb-6.2.1.post3067.dist-info}/RECORD

@@ -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.post3059.dist-info/METADATA,sha256=dGsG9jFB0KwE5eGxUB1DhoaqJmIaUzyiTDKhOJ0eSdg,4603
-nucliadb-6.2.1.post3059.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nucliadb-6.2.1.post3059.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.2.1.post3059.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.2.1.post3059.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb-6.2.1.post3059.dist-info/RECORD,,
+nucliadb-6.2.1.post3067.dist-info/METADATA,sha256=V33nguAqDHoZDoT_3Vfc2w3HTnobJN1C_uB-tlKWGMU,4603
+nucliadb-6.2.1.post3067.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nucliadb-6.2.1.post3067.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.2.1.post3067.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.2.1.post3067.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb-6.2.1.post3067.dist-info/RECORD,,