nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/search/search/query.py

@@ -18,17 +18,13 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 from nidx_protos import nodereader_pb2
 from nidx_protos.noderesources_pb2 import Resource
 
-from nucliadb.common import datamanagers
 from nucliadb.common.exceptions import InvalidQueryError
 from nucliadb.common.filter_expression import add_and_expression, parse_expression
-from nucliadb.search.search.filters import (
-    translate_label,
-)
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb_models.filters import FilterExpression
 from nucliadb_models.labels import LABEL_HIDDEN
@@ -38,7 +34,6 @@ from nucliadb_models.search import (
     SortOrder,
     SuggestOptions,
 )
-from nucliadb_protos import utils_pb2
 
 from .query_parser.old_filters import OldFilterParams, parse_old_filters
 
@@ -47,16 +42,16 @@ async def paragraph_query_to_pb(
     kbid: str,
     rid: str,
     query: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
     fields: list[str],
     filters: list[str],
     faceted: list[str],
     top_k: int,
-    range_creation_start: Optional[datetime] = None,
-    range_creation_end: Optional[datetime] = None,
-    range_modification_start: Optional[datetime] = None,
-    range_modification_end: Optional[datetime] = None,
-    sort: Optional[str] = None,
+    range_creation_start: datetime | None = None,
+    range_creation_end: datetime | None = None,
+    range_modification_start: datetime | None = None,
+    range_modification_end: datetime | None = None,
+    sort: str | None = None,
     sort_ord: str = SortOrder.DESC.value,
     with_duplicates: bool = False,
 ) -> nodereader_pb2.SearchRequest:
@@ -119,86 +114,19 @@ async def paragraph_query_to_pb(
     return request
 
 
-def expand_entities(
-    meta_cache: datamanagers.entities.EntitiesMetaCache,
-    detected_entities: list[utils_pb2.RelationNode],
-) -> list[utils_pb2.RelationNode]:
-    """
-    Iterate through duplicated entities in a kb.
-
-    The algorithm first makes it so we can look up duplicates by source and
-    by the referenced entity and expands from both directions.
-    """
-    result_entities = {entity.value: entity for entity in detected_entities}
-    duplicated_entities = meta_cache.duplicate_entities
-    duplicated_entities_by_value = meta_cache.duplicate_entities_by_value
-
-    for entity in detected_entities[:]:
-        if entity.subtype not in duplicated_entities:
-            continue
-
-        if entity.value in duplicated_entities[entity.subtype]:
-            for duplicate in duplicated_entities[entity.subtype][entity.value]:
-                result_entities[duplicate] = utils_pb2.RelationNode(
-                    ntype=utils_pb2.RelationNode.NodeType.ENTITY,
-                    subtype=entity.subtype,
-                    value=duplicate,
-                )
-
-        if entity.value in duplicated_entities_by_value[entity.subtype]:
-            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,
-                value=source_duplicate,
-            )
-
-            if source_duplicate in duplicated_entities[entity.subtype]:
-                for duplicate in duplicated_entities[entity.subtype][source_duplicate]:
-                    if duplicate == entity.value:
-                        continue
-                    result_entities[duplicate] = utils_pb2.RelationNode(
-                        ntype=utils_pb2.RelationNode.NodeType.ENTITY,
-                        subtype=entity.subtype,
-                        value=duplicate,
-                    )
-
-    return list(result_entities.values())
-
-
-def apply_entities_filter(
-    request: nodereader_pb2.SearchRequest,
-    detected_entities: list[utils_pb2.RelationNode],
-) -> list[str]:
-    added_filters = []
-    for entity_filter in [
-        f"/e/{entity.subtype}/{entity.value}"
-        for entity in detected_entities
-        if entity.ntype == utils_pb2.RelationNode.NodeType.ENTITY
-    ]:
-        if entity_filter not in added_filters:
-            added_filters.append(entity_filter)
-            # Add the entity to the filter expression (with AND)
-            entity_expr = nodereader_pb2.FilterExpression()
-            entity_expr.facet.facet = translate_label(entity_filter)
-            add_and_expression(request.field_filter, entity_expr)
-
-    return added_filters
-
-
 async def suggest_query_to_pb(
     kbid: str,
     features: list[SuggestOptions],
     query: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
     fields: list[str],
     filters: list[str],
     faceted: list[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,
-    hidden: Optional[bool] = None,
+    range_creation_start: datetime | None = None,
+    range_creation_end: datetime | None = None,
+    range_modification_start: datetime | None = None,
+    range_modification_end: datetime | None = None,
+    hidden: bool | None = None,
 ) -> nodereader_pb2.SuggestRequest:
     request = nodereader_pb2.SuggestRequest()
 
@@ -305,7 +233,7 @@ def check_supported_filters(filters: dict[str, Any], paragraph_labels: list[str]
             )
 
 
-def get_sort_field_proto(obj: SortField) -> Optional[nodereader_pb2.OrderBy.OrderField.ValueType]:
+def get_sort_field_proto(obj: SortField) -> nodereader_pb2.OrderBy.OrderField.ValueType | None:
     return {
         SortField.SCORE: None,
         SortField.CREATED: nodereader_pb2.OrderBy.OrderField.CREATED,

nucliadb/search/search/query_parser/fetcher.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-from typing import Optional, TypeVar, Union
+from typing import TypeVar
 
 from async_lru import alru_cache
 from typing_extensions import TypeIs
@@ -29,15 +29,10 @@ from nucliadb.common.maindb.utils import get_driver
 from nucliadb.search import logger
 from nucliadb.search.predict import SendToPredictError, convert_relations
 from nucliadb.search.predict_models import QueryModel
-from nucliadb.search.search.metrics import (
-    query_parse_dependency_observer,
-)
+from nucliadb.search.search.metrics import query_parse_dependency_observer
 from nucliadb.search.utilities import get_predict
 from nucliadb_models.internal.predict import QueryInfo
-from nucliadb_models.search import (
-    Image,
-    MaxTokens,
-)
+from nucliadb_models.search import Image, MaxTokens
 from nucliadb_protos import knowledgebox_pb2, utils_pb2
 
 
@@ -53,23 +48,22 @@ not_cached = NotCached()
 T = TypeVar("T")
 
 
-def is_cached(field: Union[T, NotCached]) -> TypeIs[T]:
+def is_cached(field: T | NotCached) -> TypeIs[T]:
     return not isinstance(field, NotCached)
 
 
 class FetcherCache:
-    predict_query_info: Union[Optional[QueryInfo], NotCached] = not_cached
+    predict_query_info: QueryInfo | None | NotCached = not_cached
 
     # semantic search
-    vectorset: Union[str, NotCached] = not_cached
+    vectorset: str | NotCached = not_cached
 
-    labels: Union[knowledgebox_pb2.Labels, NotCached] = not_cached
+    labels: knowledgebox_pb2.Labels | NotCached = not_cached
 
-    synonyms: Union[Optional[knowledgebox_pb2.Synonyms], NotCached] = not_cached
+    synonyms: knowledgebox_pb2.Synonyms | None | NotCached = not_cached
 
-    entities_meta_cache: Union[datamanagers.entities.EntitiesMetaCache, NotCached] = not_cached
-    deleted_entity_groups: Union[list[str], NotCached] = not_cached
-    detected_entities: Union[list[utils_pb2.RelationNode], NotCached] = not_cached
+    deleted_entity_groups: list[str] | NotCached = not_cached
+    detected_entities: list[utils_pb2.RelationNode] | NotCached = not_cached
 
 
 class Fetcher:
@@ -90,12 +84,12 @@ class Fetcher:
         kbid: str,
         *,
         query: str,
-        user_vector: Optional[list[float]],
-        vectorset: Optional[str],
+        user_vector: list[float] | None,
+        vectorset: str | None,
         rephrase: bool,
-        rephrase_prompt: Optional[str],
-        generative_model: Optional[str],
-        query_image: Optional[Image],
+        rephrase_prompt: str | None,
+        generative_model: str | None,
+        query_image: Image | None,
     ):
         self.kbid = kbid
         self.query = query
@@ -112,11 +106,11 @@ class Fetcher:
 
     # Semantic search
 
-    async def get_matryoshka_dimension(self) -> Optional[int]:
+    async def get_matryoshka_dimension(self) -> int | None:
         vectorset = await self.get_vectorset()
-        return await get_matryoshka_dimension_cached(self.kbid, vectorset)
+        return await self.get_matryoshka_dimension_cached(self.kbid, vectorset)
 
-    async def _get_user_vectorset(self) -> Optional[str]:
+    async def get_user_vectorset(self) -> str | None:
         """Returns the user's requested vectorset and validates if it does exist
         in the KB.
 
@@ -124,7 +118,7 @@ class Fetcher:
         async with self.locks.setdefault("user_vectorset", asyncio.Lock()):
             if not self.user_vectorset_validated:
                 if self.user_vectorset is not None:
-                    await validate_vectorset(self.kbid, self.user_vectorset)
+                    await self.validate_vectorset(self.kbid, self.user_vectorset)
             self.user_vectorset_validated = True
             return self.user_vectorset
 
@@ -137,7 +131,7 @@ class Fetcher:
             if is_cached(self.cache.vectorset):
                 return self.cache.vectorset
 
-            user_vectorset = await self._get_user_vectorset()
+            user_vectorset = await self.get_user_vectorset()
             if user_vectorset:
                 # user explicitly asked for a vectorset
                 self.cache.vectorset = user_vectorset
@@ -170,7 +164,7 @@ class Fetcher:
             self.cache.vectorset = vectorset
             return vectorset
 
-    async def get_query_vector(self) -> Optional[list[float]]:
+    async def get_query_vector(self) -> list[float] | None:
         if self.user_vector is not None:
             query_vector = self.user_vector
         else:
@@ -206,13 +200,20 @@ class Fetcher:
 
         return query_vector
 
-    async def get_rephrased_query(self) -> Optional[str]:
+    async def get_rephrased_query(self) -> str | None:
         query_info = await self._predict_query_endpoint()
         if query_info is None:
             return None
         return query_info.rephrased_query
 
-    async def get_semantic_min_score(self) -> Optional[float]:
+    def get_cached_rephrased_query(self) -> str | None:
+        if not is_cached(self.cache.predict_query_info):
+            return None
+        if self.cache.predict_query_info is None:
+            return None
+        return self.cache.predict_query_info.rephrased_query
+
+    async def get_semantic_min_score(self) -> float | None:
         query_info = await self._predict_query_endpoint()
         if query_info is None:
             return None
@@ -234,24 +235,6 @@ class Fetcher:
 
     # Entities
 
-    async def get_entities_meta_cache(self) -> datamanagers.entities.EntitiesMetaCache:
-        async with self.locks.setdefault("entities_meta_cache", asyncio.Lock()):
-            if is_cached(self.cache.entities_meta_cache):
-                return self.cache.entities_meta_cache
-
-            entities_meta_cache = await get_entities_meta_cache(self.kbid)
-            self.cache.entities_meta_cache = entities_meta_cache
-            return entities_meta_cache
-
-    async def get_deleted_entity_groups(self) -> list[str]:
-        async with self.locks.setdefault("deleted_entity_groups", asyncio.Lock()):
-            if is_cached(self.cache.deleted_entity_groups):
-                return self.cache.deleted_entity_groups
-
-            deleted_entity_groups = await get_deleted_entity_groups(self.kbid)
-            self.cache.deleted_entity_groups = deleted_entity_groups
-            return deleted_entity_groups
-
     async def get_detected_entities(self) -> list[utils_pb2.RelationNode]:
         async with self.locks.setdefault("detected_entities", asyncio.Lock()):
             if is_cached(self.cache.detected_entities):
@@ -275,7 +258,7 @@ class Fetcher:
 
     # Synonyms
 
-    async def get_synonyms(self) -> Optional[knowledgebox_pb2.Synonyms]:
+    async def get_synonyms(self) -> knowledgebox_pb2.Synonyms | None:
         async with self.locks.setdefault("synonyms", asyncio.Lock()):
             if is_cached(self.cache.synonyms):
                 return self.cache.synonyms
@@ -293,7 +276,7 @@ class Fetcher:
 
         return query_info.visual_llm
 
-    async def get_max_context_tokens(self, max_tokens: Optional[MaxTokens]) -> int:
+    async def get_max_context_tokens(self, max_tokens: MaxTokens | None) -> int:
         query_info = await self._predict_query_endpoint()
         if query_info is None:
             raise SendToPredictError("Error while using predict's query endpoint")
@@ -308,21 +291,21 @@ class Fetcher:
             return max_tokens.context
         return model_max
 
-    def get_max_answer_tokens(self, max_tokens: Optional[MaxTokens]) -> Optional[int]:
+    def get_max_answer_tokens(self, max_tokens: MaxTokens | None) -> int | None:
         if max_tokens is not None and max_tokens.answer is not None:
             return max_tokens.answer
         return None
 
     # Predict API
 
-    async def _predict_query_endpoint(self) -> Optional[QueryInfo]:
+    async def _predict_query_endpoint(self) -> QueryInfo | None:
         async with self.locks.setdefault("predict_query_endpoint", asyncio.Lock()):
             if is_cached(self.cache.predict_query_info):
                 return self.cache.predict_query_info
 
             # we can't call get_vectorset, as it would do a recirsive loop between
             # functions, so we'll manually parse it
-            vectorset = await self._get_user_vectorset()
+            vectorset = await self.get_user_vectorset()
             try:
                 query_info = await query_information(
                     self.kbid,
@@ -348,24 +331,28 @@ class Fetcher:
 
         return detected_entities
 
+    async def validate_vectorset(self, kbid: str, vectorset: str):
+        async with datamanagers.with_ro_transaction() as txn:
+            if not await datamanagers.vectorsets.exists(txn, kbid=kbid, vectorset_id=vectorset):
+                raise InvalidQueryError(
+                    "vectorset", f"Vectorset {vectorset} doesn't exist in your Knowledge Box"
+                )
 
-async def validate_vectorset(kbid: str, vectorset: str):
-    async with datamanagers.with_ro_transaction() as txn:
-        if not await datamanagers.vectorsets.exists(txn, kbid=kbid, vectorset_id=vectorset):
-            raise InvalidQueryError(
-                "vectorset", f"Vectorset {vectorset} doesn't exist in you Knowledge Box"
-            )
+    @alru_cache(maxsize=10)
+    async def get_matryoshka_dimension_cached(self, kbid: str, vectorset: str) -> int | None:
+        # 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": "query_information"})
 async def query_information(
     kbid: str,
     query: str,
-    semantic_model: Optional[str],
-    generative_model: Optional[str] = None,
+    semantic_model: str | None,
+    generative_model: str | None = None,
     rephrase: bool = False,
-    rephrase_prompt: Optional[str] = None,
-    query_image: Optional[Image] = None,
+    rephrase_prompt: str | None = None,
+    query_image: Image | None = None,
 ) -> QueryInfo:
     predict = get_predict()
     item = QueryModel(
@@ -385,14 +372,8 @@ async def detect_entities(kbid: str, query: str) -> list[utils_pb2.RelationNode]
     return await predict.detect_entities(kbid, query)
 
 
-@alru_cache(maxsize=None)
-async def get_matryoshka_dimension_cached(kbid: str, vectorset: 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 def get_matryoshka_dimension(kbid: str, vectorset: str | None) -> int | None:
     async with get_driver().ro_transaction() as txn:
         matryoshka_dimension = None
         if not vectorset:
@@ -414,18 +395,6 @@ async def get_classification_labels(kbid: str) -> knowledgebox_pb2.Labels:
 
 
 @query_parse_dependency_observer.wrap({"type": "synonyms"})
-async def get_kb_synonyms(kbid: str) -> Optional[knowledgebox_pb2.Synonyms]:
+async def get_kb_synonyms(kbid: str) -> knowledgebox_pb2.Synonyms | None:
     async with get_driver().ro_transaction() 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:
-    async with get_driver().ro_transaction() 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]:
-    async with get_driver().ro_transaction() as txn:
-        return list((await datamanagers.entities.get_deleted_groups(txn, kbid=kbid)).entities_groups)

nucliadb/search/search/query_parser/models.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from datetime import datetime
-from typing import Optional, Union
 
 from nidx_protos import nodereader_pb2
 from pydantic import BaseModel, ConfigDict, Field
@@ -26,6 +25,7 @@ from pydantic import BaseModel, ConfigDict, Field
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb_models import search as search_models
 from nucliadb_models.graph.requests import GraphPathQuery
+from nucliadb_models.search import MAX_RANK_FUSION_WINDOW
 from nucliadb_protos import utils_pb2
 
 ### Retrieval
@@ -46,7 +46,7 @@ KeywordQuery = _TextQuery
 
 
 class SemanticQuery(BaseModel):
-    query: Optional[list[float]]
+    query: list[float] | None
     vectorset: str
     min_score: float
 
@@ -66,11 +66,11 @@ class GraphQuery(BaseModel):
 
 
 class Query(BaseModel):
-    fulltext: Optional[FulltextQuery] = None
-    keyword: Optional[KeywordQuery] = None
-    semantic: Optional[SemanticQuery] = None
-    relation: Optional[RelationQuery] = None
-    graph: Optional[GraphQuery] = None
+    fulltext: FulltextQuery | None = None
+    keyword: KeywordQuery | None = None
+    semantic: SemanticQuery | None = None
+    relation: RelationQuery | None = None
+    graph: GraphQuery | None = None
 
 
 # filters
@@ -79,29 +79,28 @@ class Query(BaseModel):
 class Filters(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    field_expression: Optional[nodereader_pb2.FilterExpression] = None
-    paragraph_expression: Optional[nodereader_pb2.FilterExpression] = None
+    field_expression: nodereader_pb2.FilterExpression | None = None
+    paragraph_expression: nodereader_pb2.FilterExpression | None = None
     filter_expression_operator: nodereader_pb2.FilterOperator.ValueType = (
         nodereader_pb2.FilterOperator.AND
     )
 
-    autofilter: Optional[list[utils_pb2.RelationNode]] = None
     facets: list[str] = Field(default_factory=list)
-    hidden: Optional[bool] = None
-    security: Optional[search_models.RequestSecurity] = None
+    hidden: bool | None = None
+    security: search_models.RequestSecurity | None = None
     with_duplicates: bool = False
 
 
 class DateTimeFilter(BaseModel):
-    after: Optional[datetime] = None  # aka, start
-    before: Optional[datetime] = None  # aka, end
+    after: datetime | None = None  # aka, start
+    before: datetime | None = None  # aka, end
 
 
 # rank fusion
 
 
 class RankFusion(BaseModel):
-    window: int = Field(le=500)
+    window: int = Field(le=MAX_RANK_FUSION_WINDOW)
 
 
 class ReciprocalRankFusion(RankFusion):
@@ -122,7 +121,7 @@ class PredictReranker(BaseModel):
     window: int = Field(le=200)
 
 
-Reranker = Union[NoopReranker, PredictReranker]
+Reranker = NoopReranker | PredictReranker
 
 # retrieval and generation operations
 
@@ -131,8 +130,8 @@ class UnitRetrieval(BaseModel):
     query: Query
     top_k: int
     filters: Filters = Field(default_factory=Filters)
-    rank_fusion: Optional[RankFusion] = None
-    reranker: Optional[Reranker] = None
+    rank_fusion: RankFusion | None = None
+    reranker: Reranker | None = None
 
 
 # TODO: augmentation things: hydration...
@@ -141,7 +140,7 @@ class UnitRetrieval(BaseModel):
 class Generation(BaseModel):
     use_visual_llm: bool
     max_context_tokens: int
-    max_answer_tokens: Optional[int]
+    max_answer_tokens: int | None
 
 
 class ParsedQuery(BaseModel):
@@ -149,7 +148,7 @@ class ParsedQuery(BaseModel):
 
     fetcher: Fetcher
     retrieval: UnitRetrieval
-    generation: Optional[Generation] = None
+    generation: Generation | None = None
 
 
 ### Graph

nucliadb/search/search/query_parser/old_filters.py

@@ -20,7 +20,6 @@
 
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional, Union
 
 from nidx_protos.nodereader_pb2 import FilterExpression
 
@@ -36,19 +35,19 @@ from .fetcher import Fetcher
 
 @dataclass
 class OldFilterParams:
-    label_filters: Union[list[str], list[Filter]]
-    keyword_filters: Union[list[str], list[Filter]]
-    range_creation_start: Optional[datetime] = None
-    range_creation_end: Optional[datetime] = None
-    range_modification_start: Optional[datetime] = None
-    range_modification_end: Optional[datetime] = None
-    fields: Optional[list[str]] = None
-    key_filters: Optional[list[str]] = None
+    label_filters: list[str] | list[Filter]
+    keyword_filters: list[str] | list[Filter]
+    range_creation_start: datetime | None = None
+    range_creation_end: datetime | None = None
+    range_modification_start: datetime | None = None
+    range_modification_end: datetime | None = None
+    fields: list[str] | None = None
+    key_filters: list[str] | None = None
 
 
 async def parse_old_filters(
     old: OldFilterParams, fetcher: Fetcher
-) -> tuple[Optional[FilterExpression], Optional[FilterExpression]]:
+) -> tuple[FilterExpression | None, FilterExpression | None]:
     filters = []
     paragraph_filter_expression = None
 
@@ -128,6 +127,7 @@ async def parse_old_filters(
                 f.field.field_type = parts[1]
                 if len(parts) > 2:
                     f.field.field_id = parts[2]
+                expr.bool_and.operands.append(f)
             key_exprs.append(expr)
 
         if len(key_exprs) == 1:
@@ -149,8 +149,8 @@ async def parse_old_filters(
 
 
 def convert_label_filter_to_expressions(
-    fltr: Union[str, Filter], classification_labels: knowledgebox_pb2.Labels
-) -> tuple[Optional[FilterExpression], Optional[FilterExpression]]:
+    fltr: str | Filter, classification_labels: knowledgebox_pb2.Labels
+) -> tuple[FilterExpression | None, FilterExpression | None]:
     if isinstance(fltr, str):
         fltr = translate_label(fltr)
         f = FilterExpression()
@@ -174,7 +174,7 @@ def convert_label_filter_to_expressions(
 
 def split_labels(
     labels: list[str], classification_labels: knowledgebox_pb2.Labels, combinator: str, negate: bool
-) -> tuple[Optional[FilterExpression], Optional[FilterExpression]]:
+) -> tuple[FilterExpression | None, FilterExpression | None]:
     field = []
     paragraph = []
     for label in labels:
@@ -223,13 +223,14 @@ def is_paragraph_label(label: str, classification_labels: knowledgebox_pb2.Label
     if not label.startswith("/l/"):
         return False
     # Classification labels should have the form /l/labelset/label
+    # REVIEW: there's no technical reason why this has to be like this (/l/labelset could be valid)
     parts = label.split("/")
     if len(parts) < 4:
         return False
     labelset_id = parts[2]
 
     try:
-        labelset: Optional[knowledgebox_pb2.LabelSet] = classification_labels.labelset.get(labelset_id)
+        labelset: knowledgebox_pb2.LabelSet | None = classification_labels.labelset.get(labelset_id)
         if labelset is None:
             return False
         return knowledgebox_pb2.LabelSet.LabelSetKind.PARAGRAPHS in labelset.kind
@@ -238,19 +239,19 @@ def is_paragraph_label(label: str, classification_labels: knowledgebox_pb2.Label
         return False
 
 
-def convert_keyword_filter_to_expression(fltr: Union[str, Filter]) -> FilterExpression:
+def convert_keyword_filter_to_expression(fltr: str | Filter) -> FilterExpression:
     if isinstance(fltr, str):
         return convert_keyword_to_expression(fltr)
 
     f = FilterExpression()
     if fltr.all:
-        f.bool_and.operands.extend((convert_keyword_to_expression(f) for f in fltr.all))
+        f.bool_and.operands.extend(convert_keyword_to_expression(f) for f in fltr.all)
     if fltr.any:
-        f.bool_or.operands.extend((convert_keyword_to_expression(f) for f in fltr.any))
+        f.bool_or.operands.extend(convert_keyword_to_expression(f) for f in fltr.any)
     if fltr.none:
-        f.bool_not.bool_or.operands.extend((convert_keyword_to_expression(f) for f in fltr.none))
+        f.bool_not.bool_or.operands.extend(convert_keyword_to_expression(f) for f in fltr.none)
     if fltr.not_all:
-        f.bool_not.bool_and.operands.extend((convert_keyword_to_expression(f) for f in fltr.not_all))
+        f.bool_not.bool_and.operands.extend(convert_keyword_to_expression(f) for f in fltr.not_all)
 
     return f
 

nucliadb/search/search/query_parser/parsers/ask.py

@@ -17,7 +17,8 @@
 # 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/>.
 #
-from typing import Optional
+
+from typing_extensions import assert_never
 
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.models import (
@@ -26,7 +27,7 @@ from nucliadb.search.search.query_parser.models import (
 from nucliadb_models.search import AskRequest, MaxTokens
 
 
-async def parse_ask(kbid: str, item: AskRequest, *, fetcher: Optional[Fetcher] = None) -> Generation:
+async def parse_ask(kbid: str, item: AskRequest, *, fetcher: Fetcher | None = None) -> Generation:
     fetcher = fetcher or fetcher_for_ask(kbid, item)
     parser = _AskParser(kbid, item, fetcher)
     return await parser.parse()
@@ -64,9 +65,7 @@ class _AskParser:
         elif isinstance(self.item.max_tokens, MaxTokens):
             max_tokens = self.item.max_tokens
         else:  # pragma: no cover
-            # This is a trick so mypy generates an error if this branch can be reached,
-            # that is, if we are missing some ifs
-            _a: int = "a"
+            assert_never(self.item.max_tokens)
 
         max_context_tokens = await self.fetcher.get_max_context_tokens(max_tokens)
         max_answer_tokens = self.fetcher.get_max_answer_tokens(max_tokens)