nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/search/query.py

@@ -19,24 +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.ingest.orm.synonyms import Synonyms
-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 (
-    PredictVectorMissing,
-    SendToPredictError,
-    convert_relations,
-)
+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,
@@ -46,27 +41,30 @@ 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,
-    SentenceSearch,
     SortField,
-    SortFieldMap,
     SortOptions,
     SortOrder,
     SortOrderMap,
     SuggestOptions,
-    TokenSearch,
 )
 from nucliadb_models.security import RequestSecurity
 from nucliadb_protos import knowledgebox_pb2, nodereader_pb2, utils_pb2
-from nucliadb_utils import const
-from nucliadb_utils.utilities import has_feature
+from nucliadb_protos.noderesources_pb2 import Resource
 
 from .exceptions import InvalidQueryError
 
@@ -75,6 +73,8 @@ INDEX_SORTABLE_FIELDS = [
     SortField.MODIFIED,
 ]
 
+DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
+
 
 class QueryParser:
     """
@@ -86,14 +86,14 @@ class QueryParser:
     query parsing.
     """
 
-    _min_score_task: Optional[asyncio.Task] = None
     _query_information_task: Optional[asyncio.Task] = None
-    _convert_vectors_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
     _synonyms_task: Optional[asyncio.Task] = None
     _get_classification_labels_task: Optional[asyncio.Task] = None
+    _get_matryoshka_dimension_task: Optional[asyncio.Task] = None
 
     def __init__(
         self,
@@ -101,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,
@@ -121,16 +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
@@ -148,65 +160,81 @@ class QueryParser:
         self.security = security
         self.generative_model = generative_model
         self.rephrase = rephrase
-        self.query_endpoint_enabled = has_feature(
-            const.Features.PREDICT_QUERY_ENDPOINT,
-            default=False,
-            context={"kbid": self.kbid},
+        self.rephrase_prompt = rephrase_prompt
+        self.query_endpoint_used = False
+        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.SEMANTIC in self.features
+
+    @property
+    def has_relations_search(self) -> bool:
+        return SearchOptions.RELATIONS in self.features
+
+    def _get_query_information(self) -> Awaitable[QueryInfo]:
+        if self._query_information_task is None:  # pragma: no cover
+            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
         )
 
-        if len(self.filters) > 0:
-            self.filters = translate_label_filters(self.filters)
-            self.flat_filter_labels = flat_filter_labels(self.filters)
+    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
 
-    def _get_default_semantic_min_score(self) -> Awaitable[float]:
-        if self._min_score_task is None:  # pragma: no cover
-            self._min_score_task = asyncio.create_task(
-                get_default_semantic_min_score(self.kbid)
-            )
-        return self._min_score_task
+    async def _select_vectorset(self) -> Optional[str]:
+        if self.vectorset:
+            return self.vectorset
 
-    def _get_converted_vectors(self) -> Awaitable[list[float]]:
-        if self._convert_vectors_task is None:  # pragma: no cover
-            self._convert_vectors_task = asyncio.create_task(
-                convert_vectors(self.kbid, self.query)
-            )
-        return self._convert_vectors_task
+        # When vectorset is not provided we get the default from Predict API
 
-    def _get_query_information(self) -> Awaitable[QueryInfo]:
-        if self.query_endpoint_enabled is False:
-            # XXX Can be removed once query endpoint is fully enabled
-            async def static_query():
-                return QueryInfo(
-                    visual_llm=False,
-                    max_context=300_000,
-                    entities=TokenSearch(tokens=[], time=0.0),
-                    sentence=SentenceSearch(data=[], time=0.0),
-                    query=self.query,
-                )
+        try:
+            query_information = await self._get_query_information()
+        except SendToPredictError:
+            return None
 
-            return static_query()
-        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
-                )
+        if query_information.sentence is None:
+            logger.error(
+                "Asking for a vectorset but /query didn't return one", extra={"kbid": self.kbid}
             )
-        return self._query_information_task
+            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(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]]:
@@ -233,28 +261,17 @@ 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.min_score.semantic is None:
-            asyncio.ensure_future(self._get_default_semantic_min_score())
 
-        if SearchOptions.VECTOR in self.features and self.user_vector is None:
-            if self.query_endpoint_enabled:
-                asyncio.ensure_future(self._get_query_information())
-            else:
-                asyncio.ensure_future(self._get_converted_vectors())
-
-        if (SearchOptions.RELATIONS in self.features or self.autofilter) and len(
-            self.query
-        ) > 0:
-            if (
-                not self.query_endpoint_enabled
-                or SearchOptions.VECTOR not in self.features
-                or self.user_vector is not None
-            ):
-                self.query_endpoint_enabled = False
+        if self.has_vector_search and self.user_vector is None:
+            self.query_endpoint_used = True
+            asyncio.ensure_future(self._get_query_information())
+            asyncio.ensure_future(self._get_matryoshka_dimension())
+
+        if (self.has_relations_search or self.autofilter) and len(self.query) > 0:
+            if not self.query_endpoint_used:
+                # If we only need to detect entities, we don't need the query endpoint
                 asyncio.ensure_future(self._get_detected_entities())
             asyncio.ensure_future(self._get_entities_meta_cache())
             asyncio.ensure_future(self._get_deleted_entity_groups())
@@ -273,6 +290,8 @@ class QueryParser:
         request.body = self.query
         request.with_duplicates = self.with_duplicates
 
+        self.parse_sorting(request)
+
         await self._schedule_dependency_tasks()
 
         await self.parse_filters(request)
@@ -281,30 +300,29 @@ class QueryParser:
         incomplete = await self.parse_vector_search(request)
         autofilters = await self.parse_relation_search(request)
         await self.parse_synonyms(request)
-
-        self.parse_sorting(request)
-        await self.parse_min_score(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:
@@ -354,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",
@@ -369,95 +385,133 @@ 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
 
-    async def parse_min_score(self, request: nodereader_pb2.SearchRequest) -> None:
-        if self.min_score.semantic is None:
-            self.min_score.semantic = await self._get_default_semantic_min_score()
+    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()
+            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(
+                    "Vectorset unset by user or predict, using default semantic threshold",
+                    extra={"kbid": self.kbid},
+                )
+        self.min_score.semantic = semantic_min_score
         request.min_score_semantic = self.min_score.semantic
         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 SearchOptions.VECTOR not in self.features:
+        if not self.has_vector_search:
             return False
 
         node_features.inc({"type": "vectors"})
 
         incomplete = False
-        if self.vectorset is not None:
-            request.vectorset = self.vectorset
-            node_features.inc({"type": "vectorset"})
 
+        vectorset = await self._select_vectorset()
+        if vectorset is not None:
+            request.vectorset = vectorset
+
+        query_vector = None
         if self.user_vector is None:
-            if self.query_endpoint_enabled:
-                try:
-                    query_info = await self._get_query_information()
-                    if query_info and query_info.sentence:
-                        request.vector.extend(query_info.sentence.data)
-                    else:
-                        incomplete = True
-                except SendToPredictError as err:
-                    logger.warning(
-                        f"Errors on predict api trying to embedd query: {err}"
-                    )
-                    incomplete = True
-                except PredictVectorMissing:
-                    logger.warning("Predict api returned an empty vector")
-                    incomplete = True
+            try:
+                query_info = await self._get_query_information()
+            except SendToPredictError as err:
+                logger.warning(f"Errors on predict api trying to embedd query: {err}")
+                incomplete = True
             else:
-                try:
-                    request.vector.extend(await self._get_converted_vectors())
-                except SendToPredictError as err:
-                    logger.warning(
-                        f"Errors on predict api trying to embedd query: {err}"
-                    )
-                    incomplete = True
-                except PredictVectorMissing:
-                    logger.warning("Predict api returned an empty vector")
+                if query_info and query_info.sentence:
+                    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:
-            request.vector.extend(self.user_vector)
+            query_vector = self.user_vector
+
+        if query_vector is not None:
+            matryoshka_dimension = await self._get_matryoshka_dimension()
+            if matryoshka_dimension is not None:
+                # KB using a matryoshka embeddings model, cut the query vector
+                # 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 = []
-        relations_search = SearchOptions.RELATIONS in self.features
-        if relations_search or self.autofilter:
-            if not self.query_endpoint_enabled:
+        if self.has_relations_search or self.autofilter:
+            if not self.query_endpoint_used:
                 detected_entities = await self._get_detected_entities()
             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()
             detected_entities = expand_entities(meta_cache, detected_entities)
-            if relations_search:
+            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(
@@ -467,67 +521,111 @@ 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 (
-            SearchOptions.VECTOR in self.features
-            or SearchOptions.RELATIONS in self.features
-        ):
+        if self.has_vector_search or self.has_relations_search:
             raise InvalidQueryError(
                 "synonyms",
                 "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:
         return (await self._get_query_information()).visual_llm
 
-    async def get_max_context(self) -> int:
-        # Multiple by 3 is to have a good margin and guess
-        # between characters and tokens. This will be fully properly
-        # cut at the NUA API.
-        return (await self._get_query_information()).max_context * 3
+    async def get_max_tokens_context(self) -> int:
+        model_max = (await self._get_query_information()).max_context
+        if self.max_tokens is not None and self.max_tokens.context is not None:
+            if self.max_tokens.context > model_max:
+                raise InvalidQueryError(
+                    "max_tokens.context",
+                    f"Max context tokens is higher than the model's limit of {model_max}",
+                )
+            return self.max_tokens.context
+        return model_max
+
+    def get_max_tokens_answer(self) -> Optional[int]:
+        if self.max_tokens is not None and self.max_tokens.answer is not None:
+            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,
@@ -535,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)
@@ -555,44 +677,20 @@ 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
 
 
-@query_parse_dependency_observer.wrap({"type": "convert_vectors"})
-async def convert_vectors(kbid: str, query: str) -> list[float]:
-    predict = get_predict()
-    return await predict.convert_sentence_to_vector(kbid, query)
-
-
 @query_parse_dependency_observer.wrap({"type": "query_information"})
 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"})
@@ -632,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,
@@ -667,13 +763,16 @@ def parse_entities_to_filters(
         if entity_filter not in request.filter.field_labels:
             request.filter.field_labels.append(entity_filter)
             added_filters.append(entity_filter)
+
     # We need to expand the filter expression with the automatically detected entities.
     if len(added_filters) > 0:
+        # 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)
-            expanded_expression["and"].extend(expression)
-        request.filter.expression = json.dumps(expanded_expression)
+        if request.filter.labels_expression:
+            expression = json.loads(request.filter.labels_expression)
+            expanded_expression["and"].append(expression)
+        request.filter.labels_expression = json.dumps(expanded_expression)
     return added_filters
 
 
@@ -687,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()
 
@@ -696,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:
@@ -722,49 +833,28 @@ PROCESSING_STATUS_TO_PB_MAP = {
 }
 
 
-@query_parse_dependency_observer.wrap({"type": "min_score"})
-async def get_kb_model_default_min_score(kbid: str) -> Optional[float]:
-    txn = await get_read_only_transaction()
-    model = await datamanagers.kb.get_model_metadata(txn, kbid=kbid)
-    if model.HasField("default_min_score"):
-        return model.default_min_score
-    else:
-        return None
-
-
-@alru_cache(maxsize=None)
-async def get_default_semantic_min_score(kbid: str) -> float:
-    fallback = 0.7
-    model_min_score = await get_kb_model_default_min_score(kbid)
-    if model_min_score is not None:
-        return model_min_score
-    return fallback
-
-
 @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 Synonyms(txn, kbid).get()
+    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]):
@@ -783,9 +873,48 @@ 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",
             )
+
+
+@alru_cache(maxsize=None)
+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, vectorset)
+
+
+@query_parse_dependency_observer.wrap({"type": "matryoshka_dimension"})
+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]