nucliadb 6.2.1.post3201__py3-none-any.whl → 6.2.1.post3209__py3-none-any.whl

nucliadb/search/api/v1/search.py

@@ -293,7 +293,7 @@ async def search(
         hidden=await filter_hidden_resources(kbid, item.show_hidden),
         rephrase_prompt=item.rephrase_prompt,
     )
-    pb_query, incomplete_results, autofilters = await query_parser.parse()
+    pb_query, incomplete_results, autofilters, _ = await query_parser.parse()
 
     results, query_incomplete_results, queried_nodes = await node_query(
         kbid, Method.SEARCH, pb_query, target_shard_replicas=item.shards

nucliadb/search/search/find.py

@@ -106,7 +106,7 @@ async def _index_node_retrieval(
         kbid, item, generative_model=generative_model
     )
     with metrics.time("query_parse"):
-        pb_query, incomplete_results, autofilters = await query_parser.parse()
+        pb_query, incomplete_results, autofilters, rephrased_query = await query_parser.parse()
 
     with metrics.time("node_query"):
         results, query_incomplete_results, queried_nodes = await node_query(
@@ -120,6 +120,7 @@ async def _index_node_retrieval(
             results,
             kbid=kbid,
             query=pb_query.body,
+            rephrased_query=rephrased_query,
             relation_subgraph_query=pb_query.relation_subgraph,
             min_score_bm25=pb_query.min_score_bm25,
             min_score_semantic=pb_query.min_score_semantic,
@@ -194,7 +195,7 @@ async def _external_index_retrieval(
     query_parser, _, reranker = await query_parser_from_find_request(
         kbid, item, generative_model=generative_model
     )
-    search_request, incomplete_results, _ = await query_parser.parse()
+    search_request, incomplete_results, _, rephrased_query = await query_parser.parse()
 
     # Query index
     query_results = await external_index_manager.query(search_request)  # noqa
@@ -225,6 +226,7 @@ async def _external_index_retrieval(
     retrieval_results = KnowledgeboxFindResults(
         resources=find_resources,
         query=item.query,
+        rephrased_query=rephrased_query,
         total=0,
         page_number=0,
         page_size=item.top_k,
@@ -269,6 +271,7 @@ async def query_parser_from_find_request(
         kbid=kbid,
         features=item.features,
         query=item.query,
+        query_entities=item.query_entities,
         label_filters=item.filters,
         keyword_filters=item.keyword_filters,
         faceted=None,

nucliadb/search/search/find_merge.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-from typing import Iterable, Union
+from typing import Iterable, Optional, Union
 
 from nucliadb.common.external_index_providers.base import TextBlockMatch
 from nucliadb.common.ids import ParagraphId, VectorId
@@ -74,6 +74,7 @@ async def build_find_response(
     *,
     kbid: str,
     query: str,
+    rephrased_query: Optional[str],
     relation_subgraph_query: EntitiesSubgraphRequest,
     top_k: int,
     min_score_bm25: float,
@@ -143,6 +144,7 @@ async def build_find_response(
 
     find_results = KnowledgeboxFindResults(
         query=query,
+        rephrased_query=rephrased_query,
         resources=find_resources,
         best_matches=best_matches,
         relations=relations,

nucliadb/search/search/query.py

@@ -24,6 +24,7 @@ from datetime import datetime
 from typing import Any, Awaitable, Optional, Union
 
 from nucliadb.common import datamanagers
+from nucliadb.common.models_utils.from_proto import RelationNodeTypeMap
 from nucliadb.search import logger
 from nucliadb.search.predict import SendToPredictError
 from nucliadb.search.search.filters import (
@@ -49,6 +50,7 @@ from nucliadb_models.labels import LABEL_HIDDEN, translate_system_to_alias_label
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_models.search import (
     Filter,
+    KnowledgeGraphEntity,
     MaxTokens,
     MinScore,
     SearchOptions,
@@ -94,6 +96,7 @@ class QueryParser:
         keyword_filters: Union[list[str], list[Filter]],
         top_k: int,
         min_score: MinScore,
+        query_entities: Optional[list[KnowledgeGraphEntity]] = None,
         faceted: Optional[list[str]] = None,
         sort: Optional[SortOptions] = None,
         range_creation_start: Optional[datetime] = None,
@@ -120,6 +123,7 @@ class QueryParser:
         self.kbid = kbid
         self.features = features
         self.query = query
+        self.query_entities = query_entities
         self.hidden = hidden
         if self.hidden is not None:
             if self.hidden:
@@ -211,7 +215,7 @@ class QueryParser:
         if self.with_synonyms and self.query:
             asyncio.ensure_future(self.fetcher.get_synonyms())
 
-    async def parse(self) -> tuple[nodereader_pb2.SearchRequest, bool, list[str]]:
+    async def parse(self) -> tuple[nodereader_pb2.SearchRequest, bool, list[str], Optional[str]]:
         """
         :return: (request, incomplete, autofilters)
             where:
@@ -230,12 +234,13 @@ class QueryParser:
         await self.parse_filters(request)
         self.parse_document_search(request)
         self.parse_paragraph_search(request)
-        incomplete = await self.parse_vector_search(request)
+        incomplete, rephrased_query = await self.parse_vector_search(request)
+        # BUG: autofilters are not used to filter, but we say we do
         autofilters = await self.parse_relation_search(request)
         await self.parse_synonyms(request)
         await self.parse_min_score(request, incomplete)
         await self.adjust_page_size(request, self.rank_fusion, self.reranker)
-        return request, incomplete, autofilters
+        return request, incomplete, autofilters, rephrased_query
 
     async def parse_filters(self, request: nodereader_pb2.SearchRequest) -> None:
         if len(self.label_filters) > 0:
@@ -354,26 +359,41 @@ class QueryParser:
             request.paragraph = True
             node_features.inc({"type": "paragraphs"})
 
-    async def parse_vector_search(self, request: nodereader_pb2.SearchRequest) -> bool:
+    async def parse_vector_search(
+        self, request: nodereader_pb2.SearchRequest
+    ) -> tuple[bool, Optional[str]]:
         if not self.has_vector_search:
-            return False
+            return False, None
 
         node_features.inc({"type": "vectors"})
 
         vectorset = await self.fetcher.get_vectorset()
         query_vector = await self.fetcher.get_query_vector()
+        rephrased_query = await self.fetcher.get_rephrased_query()
         incomplete = query_vector is None
 
         request.vectorset = vectorset
         if query_vector is not None:
             request.vector.extend(query_vector)
 
-        return incomplete
+        return incomplete, rephrased_query
 
     async def parse_relation_search(self, request: nodereader_pb2.SearchRequest) -> list[str]:
         autofilters = []
+        # BUG: autofiler should autofilter, not enable relation search
         if self.has_relations_search or self.autofilter:
-            detected_entities = await self.fetcher.get_detected_entities()
+            if self.query_entities:
+                detected_entities = []
+                for entity in self.query_entities:
+                    relation_node = utils_pb2.RelationNode()
+                    relation_node.value = entity.name
+                    if entity.type is not None:
+                        relation_node.ntype = RelationNodeTypeMap[entity.type]
+                    if entity.subtype is not None:
+                        relation_node.subtype = entity.subtype
+                    detected_entities.append(relation_node)
+            else:
+                detected_entities = await self.fetcher.get_detected_entities()
             meta_cache = await self.fetcher.get_entities_meta_cache()
             detected_entities = expand_entities(meta_cache, detected_entities)
             if self.has_relations_search:

nucliadb/search/search/query_parser/fetcher.py

@@ -226,6 +226,12 @@ class Fetcher:
         self.cache.query_vector = query_vector
         return query_vector
 
+    async def get_rephrased_query(self) -> Optional[str]:
+        query_info = await self._predict_query_endpoint()
+        if query_info is None:
+            return None
+        return query_info.rephrased_query
+
     # Labels
 
     async def get_classification_labels(self) -> knowledgebox_pb2.Labels:

{nucliadb-6.2.1.post3201.dist-info → nucliadb-6.2.1.post3209.dist-info}/METADATA

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

{nucliadb-6.2.1.post3201.dist-info → nucliadb-6.2.1.post3209.dist-info}/RECORD

@@ -195,7 +195,7 @@ nucliadb/search/api/v1/find.py,sha256=DsnWkySu_cFajDWJIxN8DYvLL_Rm2yiCjHD8TsqPfR
 nucliadb/search/api/v1/knowledgebox.py,sha256=Hrt2h-28DDlwN3AdjMZPTYI6om0RMy9bmJvqVHvw8sE,8620
 nucliadb/search/api/v1/predict_proxy.py,sha256=QrGzo0hKjtmyGZ6pjlJHYAh4hxwVUIOTcVcerRCw7eE,3047
 nucliadb/search/api/v1/router.py,sha256=mtT07rBZcVfpa49doaw9b1tj3sdi3qLH0gn9Io6NYM0,988
-nucliadb/search/api/v1/search.py,sha256=vCj5V9kozoti0JrgU_XJhTcBucWzI4SY1B0yCSj9EQw,13638
+nucliadb/search/api/v1/search.py,sha256=aP_Iv9mi6PvmXNDX2v_t8Xhr7orD4peCY9NKo0oEnQg,13641
 nucliadb/search/api/v1/suggest.py,sha256=S0YUTAWukzZSYZJzN3T5MUgPM3599HQvG76GOCBuAbQ,5907
 nucliadb/search/api/v1/summarize.py,sha256=VAHJvE6V3xUgEBfqNKhgoxmDqCvh30RnrEIBVhMcNLU,2499
 nucliadb/search/api/v1/utils.py,sha256=5Ve-frn7LAE2jqAgB85F8RSeqxDlyA08--gS-AdOLS4,1434
@@ -210,8 +210,8 @@ nucliadb/search/search/cut.py,sha256=ytY0_GY7ocNjfxTb4aosxEp4ZfhQNDP--JkhEMGD298
 nucliadb/search/search/exceptions.py,sha256=klGLgAGGrXcSGix_W6418ZBMqDchAIGjN77ofkOScEI,1039
 nucliadb/search/search/fetch.py,sha256=XJHIFnZmXM_8Kb37lb4lg1GYG7cZ1plT-qAIb_QziX4,6184
 nucliadb/search/search/filters.py,sha256=1MkHlJjAQqoRCj7e5cEzK2HvBxGLE17I_omsjiklbtw,6476
-nucliadb/search/search/find.py,sha256=DaO3CPBQqRAw-iK_DNf_gM-aEipjtuX6oA2TbAplkxs,9901
-nucliadb/search/search/find_merge.py,sha256=5Aqz54E5GG8jw666KNncVHIJcs821ug-YwJ46YL6Br8,17363
+nucliadb/search/search/find.py,sha256=EprmlVVPHbPvcJezEJou4Msf9JlM7LD5xaVuKbk4XtY,10065
+nucliadb/search/search/find_merge.py,sha256=3FnzKFEnVemg6FO_6zveulbAU7klvsiPEBvLrpBBMg8,17450
 nucliadb/search/search/graph_strategy.py,sha256=Egcq_zn895gTUYmyQTsXj8YaUMa3HBKhcSa1GBvgzAM,31877
 nucliadb/search/search/hydrator.py,sha256=-R37gCrGxkyaiHQalnTWHNG_FCx11Zucd7qA1vQCxuw,6985
 nucliadb/search/search/merge.py,sha256=i_PTBFRqC5iTTziOMEltxLIlmokIou5hjjgR4BnoLBE,22635
@@ -219,7 +219,7 @@ nucliadb/search/search/metrics.py,sha256=81X-tahGW4n2CLvUzCPdNxNClmZqUWZjcVOGCUH
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=IaNK4dAxdXs38PoIkTdgqMDuZDjeiOtcXn3LeaT-OMw,8855
 nucliadb/search/search/predict_proxy.py,sha256=xBlh6kjuQpWRq7KsBx4pEl2PtnwljjQIiYMaTWpcCSA,3015
-nucliadb/search/search/query.py,sha256=doRdBhM928wB64v271RSyJxsRT5qd6oevImEMz4gpvw,29487
+nucliadb/search/search/query.py,sha256=AlhRw4Mick4Oab5HsKHaQpBXsVc_UUY5IpkUIwsFfU8,30577
 nucliadb/search/search/rank_fusion.py,sha256=tRGo_KlsFsVx1CQEy1iqQ6f0T1Dq1kf0axDXHuuzvvM,6946
 nucliadb/search/search/rerankers.py,sha256=0kAHES9X_FKkP7KSN9NRETFmRPKzwrFAo_54MbyvM7Q,9051
 nucliadb/search/search/shards.py,sha256=JSRSrHgHcF4sXyuZZoJdMfK0v_LHpoSRf1lCr5-K5ko,2742
@@ -233,7 +233,7 @@ nucliadb/search/search/chat/prompt.py,sha256=r2JTiRWH3YHPdeRAG5w6gD0g0fWVxdTjYIR
 nucliadb/search/search/chat/query.py,sha256=rBssR6MPSx8h2DASRMTLODaz9oGE5tNVVVeDncSrEp4,15684
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=szAOXUZ27oNY-OSa9t2hQ5HHkQQC0EX1FZz_LluJHJE,1224
-nucliadb/search/search/query_parser/fetcher.py,sha256=4ObVZSRN_dApeA2rP0yQR7l3gdvmhxY478j3pOYMssA,15528
+nucliadb/search/search/query_parser/fetcher.py,sha256=jhr__J0KmAzjdsTTadWQmD9qf6lZvqlKAfZdYjZH_UY,15742
 nucliadb/search/search/query_parser/models.py,sha256=-VlCDXUCgOroAZw1Leqhj2VMgRv_CD2w40PXXOBLaUM,2332
 nucliadb/search/search/query_parser/parser.py,sha256=JC6koS9Np1PzCfEk1Xy6mpP1HmovS_vIxxA9u-kwzos,6498
 nucliadb/standalone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -332,9 +332,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.post3201.dist-info/METADATA,sha256=m-4v4chlPV37jUjkMx61-YtplTraxGbeRj3lHOmhz_Y,4603
-nucliadb-6.2.1.post3201.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nucliadb-6.2.1.post3201.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.2.1.post3201.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.2.1.post3201.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb-6.2.1.post3201.dist-info/RECORD,,
+nucliadb-6.2.1.post3209.dist-info/METADATA,sha256=ZumCP4VHqFd-S8EEIVLceil4L-qtVSlUzwr6ao232Oo,4603
+nucliadb-6.2.1.post3209.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nucliadb-6.2.1.post3209.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.2.1.post3209.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.2.1.post3209.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb-6.2.1.post3209.dist-info/RECORD,,