nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/search/query.py

@@ -19,19 +19,19 @@
 #
 import asyncio
 import json
+import string
 from datetime import datetime
 from typing import Any, Awaitable, Optional, Union
 
 from async_lru import alru_cache
-from nucliadb_protos.noderesources_pb2 import Resource
 
 from nucliadb.common import datamanagers
-from nucliadb.middleware.transaction import get_read_only_transaction
+from nucliadb.common.maindb.utils import get_driver
 from nucliadb.search import logger
 from nucliadb.search.predict import SendToPredictError, convert_relations
 from nucliadb.search.search.filters import (
     convert_to_node_filters,
-    flat_filter_labels,
+    flatten_filter_literals,
     has_classification_label_filters,
     split_labels_by_type,
     translate_label,
@@ -41,17 +41,22 @@ from nucliadb.search.search.metrics import (
     node_features,
     query_parse_dependency_observer,
 )
+from nucliadb.search.search.rank_fusion import (
+    RankFusionAlgorithm,
+)
+from nucliadb.search.search.rerankers import (
+    Reranker,
+)
 from nucliadb.search.utilities import get_predict
-from nucliadb_models.labels import translate_system_to_alias_label
+from nucliadb_models.internal.predict import QueryInfo
+from nucliadb_models.labels import LABEL_HIDDEN, translate_system_to_alias_label
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_models.search import (
     Filter,
     MaxTokens,
     MinScore,
-    QueryInfo,
     SearchOptions,
     SortField,
-    SortFieldMap,
     SortOptions,
     SortOrder,
     SortOrderMap,
@@ -59,6 +64,7 @@ from nucliadb_models.search import (
 )
 from nucliadb_models.security import RequestSecurity
 from nucliadb_protos import knowledgebox_pb2, nodereader_pb2, utils_pb2
+from nucliadb_protos.noderesources_pb2 import Resource
 
 from .exceptions import InvalidQueryError
 
@@ -67,7 +73,6 @@ INDEX_SORTABLE_FIELDS = [
     SortField.MODIFIED,
 ]
 
-MAX_VECTOR_RESULTS_ALLOWED = 2000
 DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
 
 
@@ -82,6 +87,7 @@ class QueryParser:
     """
 
     _query_information_task: Optional[asyncio.Task] = None
+    _get_vectorset_task: Optional[asyncio.Task] = None
     _detected_entities_task: Optional[asyncio.Task] = None
     _entities_meta_cache_task: Optional[asyncio.Task] = None
     _deleted_entities_groups_task: Optional[asyncio.Task] = None
@@ -95,9 +101,9 @@ class QueryParser:
         kbid: str,
         features: list[SearchOptions],
         query: str,
-        filters: Union[list[str], list[Filter]],
-        page_number: int,
-        page_size: int,
+        label_filters: Union[list[str], list[Filter]],
+        keyword_filters: Union[list[str], list[Filter]],
+        top_k: int,
         min_score: MinScore,
         faceted: Optional[list[str]] = None,
         sort: Optional[SortOptions] = None,
@@ -107,6 +113,7 @@ class QueryParser:
         range_modification_end: Optional[datetime] = None,
         fields: Optional[list[str]] = None,
         user_vector: Optional[list[float]] = None,
+        vectorset: Optional[str] = None,
         with_duplicates: bool = False,
         with_status: Optional[ResourceProcessingStatus] = None,
         with_synonyms: bool = False,
@@ -114,17 +121,28 @@ class QueryParser:
         key_filters: Optional[list[str]] = None,
         security: Optional[RequestSecurity] = None,
         generative_model: Optional[str] = None,
-        rephrase: Optional[bool] = False,
+        rephrase: bool = False,
+        rephrase_prompt: Optional[str] = None,
         max_tokens: Optional[MaxTokens] = None,
+        hidden: Optional[bool] = None,
+        rank_fusion: Optional[RankFusionAlgorithm] = None,
+        reranker: Optional[Reranker] = None,
     ):
         self.kbid = kbid
         self.features = features
         self.query = query
-        self.filters: dict[str, Any] = convert_to_node_filters(filters)
-        self.flat_filter_labels: list[str] = []
+        self.hidden = hidden
+        if self.hidden is not None:
+            if self.hidden:
+                label_filters.append(Filter(all=[LABEL_HIDDEN]))  # type: ignore
+            else:
+                label_filters.append(Filter(none=[LABEL_HIDDEN]))  # type: ignore
+
+        self.label_filters: dict[str, Any] = convert_to_node_filters(label_filters)
+        self.flat_label_filters: list[str] = []
+        self.keyword_filters: dict[str, Any] = convert_to_node_filters(keyword_filters)
         self.faceted = faceted or []
-        self.page_number = page_number
-        self.page_size = page_size
+        self.top_k = top_k
         self.min_score = min_score
         self.sort = sort
         self.range_creation_start = range_creation_start
@@ -133,6 +151,7 @@ class QueryParser:
         self.range_modification_end = range_modification_end
         self.fields = fields or []
         self.user_vector = user_vector
+        self.vectorset = vectorset
         self.with_duplicates = with_duplicates
         self.with_status = with_status
         self.with_synonyms = with_synonyms
@@ -141,15 +160,18 @@ class QueryParser:
         self.security = security
         self.generative_model = generative_model
         self.rephrase = rephrase
+        self.rephrase_prompt = rephrase_prompt
         self.query_endpoint_used = False
-        if len(self.filters) > 0:
-            self.filters = translate_label_filters(self.filters)
-            self.flat_filter_labels = flat_filter_labels(self.filters)
+        if len(self.label_filters) > 0:
+            self.label_filters = translate_label_filters(self.label_filters)
+            self.flat_label_filters = flatten_filter_literals(self.label_filters)
         self.max_tokens = max_tokens
+        self.rank_fusion = rank_fusion
+        self.reranker = reranker
 
     @property
     def has_vector_search(self) -> bool:
-        return SearchOptions.VECTOR in self.features
+        return SearchOptions.SEMANTIC in self.features
 
     @property
     def has_relations_search(self) -> bool:
@@ -157,34 +179,62 @@ class QueryParser:
 
     def _get_query_information(self) -> Awaitable[QueryInfo]:
         if self._query_information_task is None:  # pragma: no cover
-            self._query_information_task = asyncio.create_task(
-                query_information(
-                    self.kbid, self.query, self.generative_model, self.rephrase
-                )
-            )
+            self._query_information_task = asyncio.create_task(self._query_information())
         return self._query_information_task
 
+    async def _query_information(self) -> QueryInfo:
+        vectorset = await self.select_query_vectorset()
+        return await query_information(
+            self.kbid, self.query, vectorset, self.generative_model, self.rephrase, self.rephrase_prompt
+        )
+
+    def _get_vectorset(self) -> Awaitable[Optional[str]]:
+        if self._get_vectorset_task is None:
+            self._get_vectorset_task = asyncio.create_task(self._select_vectorset())
+        return self._get_vectorset_task
+
+    async def _select_vectorset(self) -> Optional[str]:
+        if self.vectorset:
+            return self.vectorset
+
+        # When vectorset is not provided we get the default from Predict API
+
+        try:
+            query_information = await self._get_query_information()
+        except SendToPredictError:
+            return None
+
+        if query_information.sentence is None:
+            logger.error(
+                "Asking for a vectorset but /query didn't return one", extra={"kbid": self.kbid}
+            )
+            return None
+
+        for vectorset in query_information.sentence.vectors.keys():
+            self.vectorset = vectorset
+            break
+
+        return self.vectorset
+
     def _get_matryoshka_dimension(self) -> Awaitable[Optional[int]]:
         if self._get_matryoshka_dimension_task is None:
-            self._get_matryoshka_dimension_task = asyncio.create_task(
-                get_matryoshka_dimension_cached(self.kbid)
-            )
+            self._get_matryoshka_dimension_task = asyncio.create_task(self._matryoshka_dimension())
         return self._get_matryoshka_dimension_task
 
+    async def _matryoshka_dimension(self) -> Optional[int]:
+        vectorset = await self._select_vectorset()
+        return await get_matryoshka_dimension_cached(self.kbid, vectorset)
+
     def _get_detected_entities(self) -> Awaitable[list[utils_pb2.RelationNode]]:
         if self._detected_entities_task is None:  # pragma: no cover
-            self._detected_entities_task = asyncio.create_task(
-                detect_entities(self.kbid, self.query)
-            )
+            self._detected_entities_task = asyncio.create_task(detect_entities(self.kbid, self.query))
         return self._detected_entities_task
 
     def _get_entities_meta_cache(
         self,
     ) -> Awaitable[datamanagers.entities.EntitiesMetaCache]:
         if self._entities_meta_cache_task is None:
-            self._entities_meta_cache_task = asyncio.create_task(
-                get_entities_meta_cache(self.kbid)
-            )
+            self._entities_meta_cache_task = asyncio.create_task(get_entities_meta_cache(self.kbid))
         return self._entities_meta_cache_task
 
     def _get_deleted_entity_groups(self) -> Awaitable[list[str]]:
@@ -211,9 +261,7 @@ class QueryParser:
         This will schedule concurrent tasks for different data that needs to be pulled
         for the sake of the query being performed
         """
-        if len(self.filters) > 0 and has_classification_label_filters(
-            self.flat_filter_labels
-        ):
+        if len(self.label_filters) > 0 and has_classification_label_filters(self.flat_label_filters):
             asyncio.ensure_future(self._get_classification_labels())
 
         if self.has_vector_search and self.user_vector is None:
@@ -253,26 +301,28 @@ class QueryParser:
         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
 
     async def parse_filters(self, request: nodereader_pb2.SearchRequest) -> None:
-        if len(self.filters) > 0:
-            field_labels = self.flat_filter_labels
+        if len(self.label_filters) > 0:
+            field_labels = self.flat_label_filters
             paragraph_labels: list[str] = []
-            if has_classification_label_filters(self.flat_filter_labels):
+            if has_classification_label_filters(self.flat_label_filters):
                 classification_labels = await self._get_classification_labels()
                 field_labels, paragraph_labels = split_labels_by_type(
-                    self.flat_filter_labels, classification_labels
+                    self.flat_label_filters, classification_labels
                 )
-                check_supported_filters(self.filters, paragraph_labels)
+                check_supported_filters(self.label_filters, paragraph_labels)
 
             request.filter.field_labels.extend(field_labels)
             request.filter.paragraph_labels.extend(paragraph_labels)
-            request.filter.expression = json.dumps(self.filters)
+            request.filter.labels_expression = json.dumps(self.label_filters)
 
-        request.faceted.labels.extend(
-            [translate_label(facet) for facet in self.faceted]
-        )
+        if len(self.keyword_filters) > 0:
+            request.filter.keywords_expression = json.dumps(self.keyword_filters)
+
+        request.faceted.labels.extend([translate_label(facet) for facet in self.faceted])
         request.fields.extend(self.fields)
 
         if self.security is not None and len(self.security.groups) > 0:
@@ -322,9 +372,7 @@ class QueryParser:
                     order=SortOrder.DESC,
                     limit=None,
                 )
-            elif (
-                self.sort.field not in INDEX_SORTABLE_FIELDS and self.sort.limit is None
-            ):
+            elif self.sort.field not in INDEX_SORTABLE_FIELDS and self.sort.limit is None:
                 raise InvalidQueryError(
                     "sort_field",
                     f"Sort by '{self.sort.field}' requires setting a sort limit",
@@ -337,35 +385,32 @@ class QueryParser:
             # have consistent results, we must limit them
             request.result_per_page = self.sort.limit
         else:
-            request.result_per_page = self.page_number * self.page_size + self.page_size
+            request.result_per_page = self.top_k
 
-        sort_field = SortFieldMap[self.sort.field] if self.sort else None
+        sort_field = get_sort_field_proto(self.sort.field) if self.sort else None
         if sort_field is not None:
             request.order.sort_by = sort_field
             request.order.type = SortOrderMap[self.sort.order]  # type: ignore
 
-        if (
-            self.has_vector_search
-            and request.result_per_page > MAX_VECTOR_RESULTS_ALLOWED
-        ):
-            raise InvalidQueryError(
-                "page_size",
-                f"Pagination of semantic results limit reached: {MAX_VECTOR_RESULTS_ALLOWED}. If you want to paginate through all results, please disable the vector search feature.",  # noqa: E501
-            )
-
-    async def parse_min_score(
-        self, request: nodereader_pb2.SearchRequest, incomplete: bool
-    ) -> None:
+    async def parse_min_score(self, request: nodereader_pb2.SearchRequest, incomplete: bool) -> None:
         semantic_min_score = DEFAULT_GENERIC_SEMANTIC_THRESHOLD
         if self.min_score.semantic is not None:
             semantic_min_score = self.min_score.semantic
         elif self.has_vector_search and not incomplete:
             query_information = await self._get_query_information()
-            if query_information.semantic_threshold is not None:
-                semantic_min_score = query_information.semantic_threshold
+            vectorset = await self._select_vectorset()
+            if vectorset is not None:
+                semantic_threshold = query_information.semantic_thresholds.get(vectorset, None)
+                if semantic_threshold is not None:
+                    semantic_min_score = semantic_threshold
+                else:
+                    logger.warning(
+                        "Semantic threshold not found in query information, using default",
+                        extra={"kbid": self.kbid},
+                    )
             else:
                 logger.warning(
-                    "Semantic threshold not found in query information, using default",
+                    "Vectorset unset by user or predict, using default semantic threshold",
                     extra={"kbid": self.kbid},
                 )
         self.min_score.semantic = semantic_min_score
@@ -373,15 +418,34 @@ class QueryParser:
         request.min_score_bm25 = self.min_score.bm25
 
     def parse_document_search(self, request: nodereader_pb2.SearchRequest) -> None:
-        if SearchOptions.DOCUMENT in self.features:
+        if SearchOptions.FULLTEXT in self.features:
             request.document = True
             node_features.inc({"type": "documents"})
 
     def parse_paragraph_search(self, request: nodereader_pb2.SearchRequest) -> None:
-        if SearchOptions.PARAGRAPH in self.features:
+        if SearchOptions.KEYWORD in self.features:
             request.paragraph = True
             node_features.inc({"type": "paragraphs"})
 
+    async def select_query_vectorset(self) -> Optional[str]:
+        """Set and return the requested vectorset parameter (if used) validated
+        for the current KB.
+
+        """
+        if not self.vectorset:
+            return None
+
+        # validate vectorset
+        async with datamanagers.with_ro_transaction() as txn:
+            if not await datamanagers.vectorsets.exists(
+                txn, kbid=self.kbid, vectorset_id=self.vectorset
+            ):
+                raise InvalidQueryError(
+                    "vectorset",
+                    f"Vectorset {self.vectorset} doesn't exist in you Knowledge Box",
+                )
+        return self.vectorset
+
     async def parse_vector_search(self, request: nodereader_pb2.SearchRequest) -> bool:
         if not self.has_vector_search:
             return False
@@ -389,6 +453,11 @@ class QueryParser:
         node_features.inc({"type": "vectors"})
 
         incomplete = False
+
+        vectorset = await self._select_vectorset()
+        if vectorset is not None:
+            request.vectorset = vectorset
+
         query_vector = None
         if self.user_vector is None:
             try:
@@ -398,11 +467,24 @@ class QueryParser:
                 incomplete = True
             else:
                 if query_info and query_info.sentence:
-                    query_vector = query_info.sentence.data
+                    if vectorset:
+                        if vectorset in query_info.sentence.vectors:
+                            query_vector = query_info.sentence.vectors[vectorset]
+                        else:
+                            incomplete = True
+                    else:
+                        for vectorset_id, vector in query_info.sentence.vectors.items():
+                            if vector:
+                                query_vector = vector
+                                break
+                        else:
+                            incomplete = True
+
                 else:
                     incomplete = True
         else:
             query_vector = self.user_vector
+
         if query_vector is not None:
             matryoshka_dimension = await self._get_matryoshka_dimension()
             if matryoshka_dimension is not None:
@@ -410,11 +492,10 @@ class QueryParser:
                 # accordingly
                 query_vector = query_vector[:matryoshka_dimension]
             request.vector.extend(query_vector)
+
         return incomplete
 
-    async def parse_relation_search(
-        self, request: nodereader_pb2.SearchRequest
-    ) -> list[str]:
+    async def parse_relation_search(self, request: nodereader_pb2.SearchRequest) -> list[str]:
         autofilters = []
         if self.has_relations_search or self.autofilter:
             if not self.query_endpoint_used:
@@ -422,9 +503,7 @@ class QueryParser:
             else:
                 query_info_result = await self._get_query_information()
                 if query_info_result.entities:
-                    detected_entities = convert_relations(
-                        query_info_result.entities.dict()
-                    )
+                    detected_entities = convert_relations(query_info_result.entities.model_dump())
                 else:
                     detected_entities = []
             meta_cache = await self._get_entities_meta_cache()
@@ -432,9 +511,7 @@ class QueryParser:
             if self.has_relations_search:
                 request.relation_subgraph.entry_points.extend(detected_entities)
                 request.relation_subgraph.depth = 1
-                request.relation_subgraph.deleted_groups.extend(
-                    await self._get_deleted_entity_groups()
-                )
+                request.relation_subgraph.deleted_groups.extend(await self._get_deleted_entity_groups())
                 for group_id, deleted_entities in meta_cache.deleted_entities.items():
                     request.relation_subgraph.deleted_entities.append(
                         nodereader_pb2.EntitiesSubgraphRequest.DeletedEntities(
@@ -444,13 +521,21 @@ class QueryParser:
                 node_features.inc({"type": "relations"})
             if self.autofilter:
                 entity_filters = parse_entities_to_filters(request, detected_entities)
-                autofilters.extend(
-                    [translate_system_to_alias_label(e) for e in entity_filters]
-                )
+                autofilters.extend([translate_system_to_alias_label(e) for e in entity_filters])
         return autofilters
 
     async def parse_synonyms(self, request: nodereader_pb2.SearchRequest) -> None:
-        if not self.with_synonyms:
+        """
+        Replace the terms in the query with an expression that will make it match with the configured synonyms.
+        We're using the Tantivy's query language here: https://docs.rs/tantivy/latest/tantivy/query/struct.QueryParser.html
+
+        Example:
+        - Synonyms: Foo -> Bar, Baz
+        - Query: "What is Foo?"
+        - Advanced Query: "What is (Foo OR Bar OR Baz)?"
+        """
+        if not self.with_synonyms or not self.query:
+            # Nothing to do
             return
 
         if self.has_vector_search or self.has_relations_search:
@@ -459,27 +544,32 @@ class QueryParser:
                 "Search with custom synonyms is only supported on paragraph and document search",
             )
 
-        if not self.query:
-            # Nothing to do
-            return
-
         synonyms = await self._get_synomyns()
         if synonyms is None:
             # No synonyms found
             return
 
-        synonyms_found: list[str] = []
-        advanced_query = []
-        for term in self.query.split(" "):
-            advanced_query.append(term)
-            term_synonyms = synonyms.terms.get(term)
-            if term_synonyms is None or len(term_synonyms.synonyms) == 0:
-                # No synonyms found for this term
-                continue
-            synonyms_found.extend(term_synonyms.synonyms)
-
-        if len(synonyms_found):
-            request.advanced_query = " OR ".join(advanced_query + synonyms_found)
+        # Calculate term variants: 'term' -> '(term OR synonym1 OR synonym2)'
+        variants: dict[str, str] = {}
+        for term, term_synonyms in synonyms.terms.items():
+            if len(term_synonyms.synonyms) > 0:
+                variants[term] = "({})".format(" OR ".join([term] + list(term_synonyms.synonyms)))
+
+        # Split the query into terms
+        query_terms = self.query.split()
+
+        # Remove punctuation from the query terms
+        clean_query_terms = [term.strip(string.punctuation) for term in query_terms]
+
+        # Replace the original terms with the variants if the cleaned term is in the variants
+        term_with_synonyms_found = False
+        for index, clean_term in enumerate(clean_query_terms):
+            if clean_term in variants:
+                term_with_synonyms_found = True
+                query_terms[index] = query_terms[index].replace(clean_term, variants[clean_term])
+
+        if term_with_synonyms_found:
+            request.advanced_query = " ".join(query_terms)
             request.ClearField("body")
 
     async def get_visual_llm_enabled(self) -> bool:
@@ -501,17 +591,41 @@ class QueryParser:
             return self.max_tokens.answer
         return None
 
+    async def adjust_page_size(
+        self,
+        request: nodereader_pb2.SearchRequest,
+        rank_fusion: Optional[RankFusionAlgorithm],
+        reranker: Optional[Reranker],
+    ):
+        """Adjust requested page size depending on rank fusion and reranking algorithms.
+
+        Some rerankers want more results than the requested by the user so
+        reranking can have more choices.
+
+        """
+        rank_fusion_window = 0
+        if rank_fusion is not None:
+            rank_fusion_window = rank_fusion.window
+
+        reranker_window = 0
+        if reranker is not None:
+            reranker_window = reranker.window or 0
+
+        request.result_per_page = max(
+            request.result_per_page,
+            rank_fusion_window,
+            reranker_window,
+        )
+
 
 async def paragraph_query_to_pb(
     kbid: str,
-    features: list[SearchOptions],
     rid: str,
     query: str,
     fields: list[str],
     filters: list[str],
     faceted: list[str],
-    page_number: int,
-    page_size: int,
+    top_k: int,
     range_creation_start: Optional[datetime] = None,
     range_creation_end: Optional[datetime] = None,
     range_modification_start: Optional[datetime] = None,
@@ -519,13 +633,37 @@ async def paragraph_query_to_pb(
     sort: Optional[str] = None,
     sort_ord: str = SortOrder.DESC.value,
     with_duplicates: bool = False,
-) -> nodereader_pb2.ParagraphSearchRequest:
-    request = nodereader_pb2.ParagraphSearchRequest()
-    request.with_duplicates = with_duplicates
+) -> nodereader_pb2.SearchRequest:
+    request = nodereader_pb2.SearchRequest()
+    request.paragraph = True
 
     # We need to ask for all and cut later
     request.page_number = 0
-    request.result_per_page = page_number * page_size + page_size
+    request.result_per_page = top_k
+
+    request.body = query
+
+    # we don't have a specific filter only for resource_ids but key_filters
+    # parse "rid" and "rid/field" like ids, so it does the job
+    request.key_filters.append(rid)
+
+    if len(filters) > 0:
+        field_labels = filters
+        paragraph_labels: list[str] = []
+        if has_classification_label_filters(filters):
+            classification_labels = await get_classification_labels(kbid)
+            field_labels, paragraph_labels = split_labels_by_type(filters, classification_labels)
+        request.filter.field_labels.extend(field_labels)
+        request.filter.paragraph_labels.extend(paragraph_labels)
+
+    request.faceted.labels.extend([translate_label(facet) for facet in faceted])
+    request.fields.extend(fields)
+
+    if sort:
+        request.order.field = sort
+        request.order.type = sort_ord  # type: ignore
+
+    request.with_duplicates = with_duplicates
 
     if range_creation_start is not None:
         request.timestamps.from_created.FromDatetime(range_creation_start)
@@ -539,26 +677,6 @@ async def paragraph_query_to_pb(
     if range_modification_end is not None:
         request.timestamps.to_modified.FromDatetime(range_modification_end)
 
-    if SearchOptions.PARAGRAPH in features:
-        request.uuid = rid
-        request.body = query
-        if len(filters) > 0:
-            field_labels = filters
-            paragraph_labels: list[str] = []
-            if has_classification_label_filters(filters):
-                classification_labels = await get_classification_labels(kbid)
-                field_labels, paragraph_labels = split_labels_by_type(
-                    filters, classification_labels
-                )
-            request.filter.field_labels.extend(field_labels)
-            request.filter.paragraph_labels.extend(paragraph_labels)
-
-        request.faceted.labels.extend([translate_label(facet) for facet in faceted])
-        if sort:
-            request.order.field = sort
-            request.order.type = sort_ord  # type: ignore
-        request.fields.extend(fields)
-
     return request
 
 
@@ -566,11 +684,13 @@ async def paragraph_query_to_pb(
 async def query_information(
     kbid: str,
     query: str,
+    semantic_model: Optional[str],
     generative_model: Optional[str] = None,
     rephrase: bool = False,
+    rephrase_prompt: Optional[str] = None,
 ) -> QueryInfo:
     predict = get_predict()
-    return await predict.query(kbid, query, generative_model, rephrase)
+    return await predict.query(kbid, query, semantic_model, generative_model, rephrase, rephrase_prompt)
 
 
 @query_parse_dependency_observer.wrap({"type": "detect_entities"})
@@ -610,9 +730,7 @@ def expand_entities(
                 )
 
         if entity.value in duplicated_entities_by_value[entity.subtype]:
-            source_duplicate = duplicated_entities_by_value[entity.subtype][
-                entity.value
-            ]
+            source_duplicate = duplicated_entities_by_value[entity.subtype][entity.value]
             result_entities[source_duplicate] = utils_pb2.RelationNode(
                 ntype=utils_pb2.RelationNode.NodeType.ENTITY,
                 subtype=entity.subtype,
@@ -651,10 +769,10 @@ def parse_entities_to_filters(
         # So far, autofilters feature will only yield 'and' expressions with the detected entities.
         # More complex autofilters can be added here if we leverage the query endpoint.
         expanded_expression = {"and": [{"literal": entity} for entity in added_filters]}
-        if request.filter.expression:
-            expression = json.loads(request.filter.expression)
+        if request.filter.labels_expression:
+            expression = json.loads(request.filter.labels_expression)
             expanded_expression["and"].append(expression)
-        request.filter.expression = json.dumps(expanded_expression)
+        request.filter.labels_expression = json.dumps(expanded_expression)
     return added_filters
 
 
@@ -668,6 +786,7 @@ def suggest_query_to_pb(
     range_creation_end: Optional[datetime] = None,
     range_modification_start: Optional[datetime] = None,
     range_modification_end: Optional[datetime] = None,
+    hidden: Optional[bool] = None,
 ) -> nodereader_pb2.SuggestRequest:
     request = nodereader_pb2.SuggestRequest()
 
@@ -677,10 +796,21 @@ def suggest_query_to_pb(
 
     if SuggestOptions.PARAGRAPH in features:
         request.features.append(nodereader_pb2.SuggestFeatures.PARAGRAPHS)
-        filters = [translate_label(fltr) for fltr in filters]
-        request.filter.field_labels.extend(filters)
         request.fields.extend(fields)
 
+        if hidden is not None:
+            if hidden:
+                filters.append(Filter(all=[LABEL_HIDDEN]))  # type: ignore
+            else:
+                filters.append(Filter(none=[LABEL_HIDDEN]))  # type: ignore
+
+        expression = convert_to_node_filters(filters)
+        if expression:
+            expression = translate_label_filters(expression)
+
+        request.filter.field_labels.extend(flatten_filter_literals(expression))
+        request.filter.labels_expression = json.dumps(expression)
+
     if range_creation_start is not None:
         request.timestamps.from_created.FromDatetime(range_creation_start)
     if range_creation_end is not None:
@@ -705,28 +835,26 @@ PROCESSING_STATUS_TO_PB_MAP = {
 
 @query_parse_dependency_observer.wrap({"type": "synonyms"})
 async def get_kb_synonyms(kbid: str) -> Optional[knowledgebox_pb2.Synonyms]:
-    txn = await get_read_only_transaction()
-    return await datamanagers.synonyms.get(txn, kbid=kbid)
+    async with get_driver().transaction(read_only=True) as txn:
+        return await datamanagers.synonyms.get(txn, kbid=kbid)
 
 
 @query_parse_dependency_observer.wrap({"type": "entities_meta_cache"})
 async def get_entities_meta_cache(kbid: str) -> datamanagers.entities.EntitiesMetaCache:
-    txn = await get_read_only_transaction()
-    return await datamanagers.entities.get_entities_meta_cache(txn, kbid=kbid)
+    async with get_driver().transaction(read_only=True) as txn:
+        return await datamanagers.entities.get_entities_meta_cache(txn, kbid=kbid)
 
 
 @query_parse_dependency_observer.wrap({"type": "deleted_entities_groups"})
 async def get_deleted_entity_groups(kbid: str) -> list[str]:
-    txn = await get_read_only_transaction()
-    return list(
-        (await datamanagers.entities.get_deleted_groups(txn, kbid=kbid)).entities_groups
-    )
+    async with get_driver().transaction(read_only=True) as txn:
+        return list((await datamanagers.entities.get_deleted_groups(txn, kbid=kbid)).entities_groups)
 
 
 @query_parse_dependency_observer.wrap({"type": "classification_labels"})
 async def get_classification_labels(kbid: str) -> knowledgebox_pb2.Labels:
-    txn = await get_read_only_transaction()
-    return await datamanagers.labels.get_labels(txn, kbid=kbid)
+    async with get_driver().transaction(read_only=True) as txn:
+        return await datamanagers.labels.get_labels(txn, kbid=kbid)
 
 
 def check_supported_filters(filters: dict[str, Any], paragraph_labels: list[str]):
@@ -745,8 +873,16 @@ def check_supported_filters(filters: dict[str, Any], paragraph_labels: list[str]
             "Paragraph labels can only be used with 'all' filter",
         )
     for term in filters["and"]:
-        # Nested expressions are not allowed with paragraph labels
-        if "literal" not in term:
+        # Nested expressions are not allowed with paragraph labels (only "literal" and "not(literal)")
+        if "not" in term:
+            subterm = term["not"]
+            if "literal" not in subterm:
+                # AND (NOT( X )) where X is anything other than a literal
+                raise InvalidQueryError(
+                    "filters",
+                    "Paragraph labels can only be used with 'all' filter",
+                )
+        elif "literal" not in term:
             raise InvalidQueryError(
                 "filters",
                 "Paragraph labels can only be used with 'all' filter",
@@ -754,12 +890,31 @@ def check_supported_filters(filters: dict[str, Any], paragraph_labels: list[str]
 
 
 @alru_cache(maxsize=None)
-async def get_matryoshka_dimension_cached(kbid: str) -> Optional[int]:
+async def get_matryoshka_dimension_cached(kbid: str, vectorset: Optional[str]) -> Optional[int]:
     # This can be safely cached as the matryoshka dimension is not expected to change
-    return await get_matryoshka_dimension(kbid)
+    return await get_matryoshka_dimension(kbid, vectorset)
 
 
 @query_parse_dependency_observer.wrap({"type": "matryoshka_dimension"})
-async def get_matryoshka_dimension(kbid: str) -> Optional[int]:
-    txn = await get_read_only_transaction()
-    return await datamanagers.kb.get_matryoshka_vector_dimension(txn, kbid=kbid)
+async def get_matryoshka_dimension(kbid: str, vectorset: Optional[str]) -> Optional[int]:
+    async with get_driver().transaction(read_only=True) as txn:
+        matryoshka_dimension = None
+        if not vectorset:
+            # XXX this should be migrated once we remove the "default" vectorset
+            # concept
+            matryoshka_dimension = await datamanagers.kb.get_matryoshka_vector_dimension(txn, kbid=kbid)
+        else:
+            vectorset_config = await datamanagers.vectorsets.get(txn, kbid=kbid, vectorset_id=vectorset)
+            if vectorset_config is not None and vectorset_config.vectorset_index_config.vector_dimension:
+                matryoshka_dimension = vectorset_config.vectorset_index_config.vector_dimension
+
+        return matryoshka_dimension
+
+
+def get_sort_field_proto(obj: SortField) -> Optional[nodereader_pb2.OrderBy.OrderField.ValueType]:
+    return {
+        SortField.SCORE: None,
+        SortField.CREATED: nodereader_pb2.OrderBy.OrderField.CREATED,
+        SortField.MODIFIED: nodereader_pb2.OrderBy.OrderField.MODIFIED,
+        SortField.TITLE: None,
+    }[obj]