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

nucliadb/search/search/query_parser/exceptions.py

@@ -0,0 +1,22 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+
+
+class ParserError(ValueError): ...

nucliadb/search/search/query_parser/models.py

@@ -0,0 +1,101 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Optional
+
+from pydantic import (
+    BaseModel,
+    Field,
+)
+
+from nucliadb_models import search as search_models
+
+### Retrieval
+
+# filters
+
+
+class DateTimeFilter(BaseModel):
+    after: Optional[datetime] = None  # aka, start
+    before: Optional[datetime] = None  # aka, end
+
+
+# rank fusion
+
+
+class RankFusion(BaseModel):
+    window: int = Field(le=500)
+
+
+class ReciprocalRankFusion(RankFusion):
+    k: float = Field(default=60.0)
+    boosting: search_models.ReciprocalRankFusionWeights = Field(
+        default_factory=search_models.ReciprocalRankFusionWeights
+    )
+
+
+# reranking
+
+
+class Reranker(BaseModel): ...
+
+
+class NoopReranker(Reranker): ...
+
+
+class MultiMatchBoosterReranker(Reranker): ...
+
+
+class PredictReranker(Reranker):
+    window: int = Field(le=200)
+
+
+# retrieval operation
+
+
+@dataclass
+class UnitRetrieval:
+    top_k: int
+    rank_fusion: RankFusion
+    reranker: Reranker
+
+
+### Catalog
+
+
+class CatalogFilters(BaseModel):
+    labels: dict[str, Any] = Field(
+        default_factory=dict, description="Labels filter expression, like, `{and: {not: ...}, ...}`"
+    )
+    creation: DateTimeFilter
+    modification: DateTimeFilter
+    with_status: Optional[search_models.ResourceProcessingStatus] = None
+
+
+class CatalogQuery(BaseModel):
+    kbid: str
+    query: str
+    filters: CatalogFilters
+    sort: search_models.SortOptions
+    faceted: list[str]
+    page_size: int
+    page_number: int

nucliadb/search/search/query_parser/parser.py

@@ -0,0 +1,183 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 Any
+
+from pydantic import ValidationError
+
+from nucliadb.search.search.filters import (
+    convert_to_node_filters,
+    translate_label_filters,
+)
+from nucliadb.search.search.query_parser.exceptions import ParserError
+from nucliadb.search.search.query_parser.models import (
+    CatalogFilters,
+    CatalogQuery,
+    DateTimeFilter,
+    MultiMatchBoosterReranker,
+    NoopReranker,
+    PredictReranker,
+    RankFusion,
+    ReciprocalRankFusion,
+    Reranker,
+    UnitRetrieval,
+)
+from nucliadb_models import search as search_models
+from nucliadb_models.labels import LABEL_HIDDEN
+from nucliadb_models.search import (
+    Filter,
+    FindRequest,
+    SortField,
+    SortOptions,
+    SortOrder,
+)
+
+
+def parse_find(item: FindRequest) -> UnitRetrieval:
+    parser = _FindParser(item)
+    return parser.parse()
+
+
+class _FindParser:
+    def __init__(self, item: FindRequest):
+        self.item = item
+
+    def parse(self) -> UnitRetrieval:
+        top_k = self._parse_top_k()
+        try:
+            rank_fusion = self._parse_rank_fusion()
+        except ValidationError as exc:
+            raise ParserError(f"Parsing error in rank fusion: {str(exc)}") from exc
+        try:
+            reranker = self._parse_reranker()
+        except ValidationError as exc:
+            raise ParserError(f"Parsing error in reranker: {str(exc)}") from exc
+
+        # Adjust retrieval windows. Our current implementation assume:
+        # `top_k <= reranker.window <= rank_fusion.window`
+        # and as rank fusion is done before reranking, we must ensure rank
+        # fusion window is at least, the reranker window
+        if isinstance(reranker, PredictReranker):
+            rank_fusion.window = max(rank_fusion.window, reranker.window)
+
+        return UnitRetrieval(
+            top_k=top_k,
+            rank_fusion=rank_fusion,
+            reranker=reranker,
+        )
+
+    def _parse_top_k(self) -> int:
+        assert self.item.top_k is not None, "top_k must have an int value"
+        top_k = self.item.top_k
+        return top_k
+
+    def _parse_rank_fusion(self) -> RankFusion:
+        rank_fusion: RankFusion
+
+        top_k = self._parse_top_k()
+        window = min(top_k, 500)
+
+        if isinstance(self.item.rank_fusion, search_models.RankFusionName):
+            if self.item.rank_fusion == search_models.RankFusionName.RECIPROCAL_RANK_FUSION:
+                rank_fusion = ReciprocalRankFusion(window=window)
+            else:
+                raise ParserError(f"Unknown rank fusion algorithm: {self.item.rank_fusion}")
+
+        elif isinstance(self.item.rank_fusion, search_models.ReciprocalRankFusion):
+            user_window = self.item.rank_fusion.window
+            rank_fusion = ReciprocalRankFusion(
+                k=self.item.rank_fusion.k,
+                boosting=self.item.rank_fusion.boosting,
+                window=min(max(user_window or 0, top_k), 500),
+            )
+
+        else:
+            raise ParserError(f"Unknown rank fusion {self.item.rank_fusion}")
+
+        return rank_fusion
+
+    def _parse_reranker(self) -> Reranker:
+        reranking: Reranker
+
+        top_k = self._parse_top_k()
+
+        if isinstance(self.item.reranker, search_models.RerankerName):
+            if self.item.reranker == search_models.RerankerName.NOOP:
+                reranking = NoopReranker()
+
+            elif self.item.reranker == search_models.RerankerName.MULTI_MATCH_BOOSTER:
+                reranking = MultiMatchBoosterReranker()
+
+            elif self.item.reranker == search_models.RerankerName.PREDICT_RERANKER:
+                # for predict rearnker, by default, we want a x2 factor with a
+                # top of 200 results
+                reranking = PredictReranker(window=min(top_k * 2, 200))
+
+            else:
+                raise ParserError(f"Unknown reranker algorithm: {self.item.reranker}")
+
+        elif isinstance(self.item.reranker, search_models.PredictReranker):
+            user_window = self.item.reranker.window
+            reranking = PredictReranker(window=min(max(user_window or 0, top_k), 200))
+
+        else:
+            raise ParserError(f"Unknown reranker {self.item.reranker}")
+
+        return reranking
+
+
+def parse_catalog(kbid: str, item: search_models.CatalogRequest) -> CatalogQuery:
+    if item.hidden:
+        hidden_filter = Filter(all=[LABEL_HIDDEN])
+    else:
+        hidden_filter = Filter(none=[LABEL_HIDDEN])
+    label_filters: dict[str, Any] = convert_to_node_filters(item.filters + [hidden_filter])  # type: ignore
+    if len(label_filters) > 0:
+        label_filters = translate_label_filters(label_filters)
+
+    sort = item.sort
+    if sort is None:
+        # By default we sort by creation date (most recent first)
+        sort = SortOptions(
+            field=SortField.CREATED,
+            order=SortOrder.DESC,
+            limit=None,
+        )
+
+    return CatalogQuery(
+        kbid=kbid,
+        query=item.query,
+        filters=CatalogFilters(
+            labels=label_filters,
+            creation=DateTimeFilter(
+                after=item.range_creation_start,
+                before=item.range_creation_end,
+            ),
+            modification=DateTimeFilter(
+                after=item.range_modification_start,
+                before=item.range_modification_end,
+            ),
+            with_status=item.with_status,
+        ),
+        sort=sort,
+        faceted=item.faceted,
+        page_number=item.page_number,
+        page_size=item.page_size,
+    )

nucliadb/search/search/rank_fusion.py

@@ -0,0 +1,204 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+import logging
+from abc import ABC, abstractmethod
+from typing import Iterable
+
+from nucliadb.common.external_index_providers.base import TextBlockMatch
+from nucliadb.common.ids import ParagraphId
+from nucliadb.search.search.query_parser import models as parser_models
+from nucliadb_models.search import SCORE_TYPE
+from nucliadb_telemetry.metrics import Observer
+
+logger = logging.getLogger(__name__)
+
+rank_fusion_observer = Observer(
+    "rank_fusion",
+    labels={"type": ""},
+    buckets=[
+        0.001,
+        0.0025,
+        0.005,
+        0.01,
+        0.025,
+        0.05,
+        0.1,
+        0.25,
+        0.5,
+        1.0,
+    ],
+)
+
+
+class RankFusionAlgorithm(ABC):
+    def __init__(self, window: int):
+        self._window = window
+
+    @property
+    def window(self) -> int:
+        """Phony number used to compute the number of elements to retrieve and
+        feed the rank fusion algorithm.
+
+        This is here for convinience, but a query plan should be the way to go.
+
+        """
+        return self._window
+
+    def fuse(
+        self, keyword: Iterable[TextBlockMatch], semantic: Iterable[TextBlockMatch]
+    ) -> list[TextBlockMatch]:
+        """Fuse keyword and semantic results and return a list with the merged
+        results.
+
+        """
+        merged = self._fuse(keyword, semantic)
+        return merged
+
+    @abstractmethod
+    def _fuse(
+        self, keyword: Iterable[TextBlockMatch], semantic: Iterable[TextBlockMatch]
+    ) -> list[TextBlockMatch]: ...
+
+
+class LegacyRankFusion(RankFusionAlgorithm):
+    """Legacy algorithm that given results from keyword and semantic search,
+    mixes them in the following way:
+    - 1st result from keyword search
+    - 2nd result from semantic search
+    - 2 keyword results and 1 semantic (and repeat)
+
+    """
+
+    @rank_fusion_observer.wrap({"type": "legacy"})
+    def _fuse(
+        self, keyword: Iterable[TextBlockMatch], semantic: Iterable[TextBlockMatch]
+    ) -> list[TextBlockMatch]:
+        merged: list[TextBlockMatch] = []
+
+        # sort results by it's score before merging them
+        keyword = [k for k in sorted(keyword, key=lambda r: r.score, reverse=True)]
+        semantic = [s for s in sorted(semantic, key=lambda r: r.score, reverse=True)]
+
+        for k in keyword:
+            merged.append(k)
+
+        nextpos = 1
+        for s in semantic:
+            merged.insert(nextpos, s)
+            nextpos += 3
+
+        return merged
+
+
+class ReciprocalRankFusion(RankFusionAlgorithm):
+    """Rank-based rank fusion algorithm. Discounts the weight of documents
+    occurring deep in retrieved lists using a reciprocal distribution. It can be
+    parametrized with weights to boost retrievers.
+
+    RRF = Σ(r ∈ R) (1 / (k + r(d)) · w(r))
+
+    where:
+    - d is a document
+    - R is the set of retrievers
+    - k (constant)
+    - r(d) rank of document d in reranker r
+    - w(r) weight (boost) for retriever r
+
+    RRF boosts matches from multiple retrievers and deduplicate them
+
+    """
+
+    def __init__(
+        self,
+        k: float = 60.0,
+        *,
+        window: int,
+        keyword_weight: float = 1.0,
+        semantic_weight: float = 1.0,
+    ):
+        super().__init__(window)
+        # Constant used in RRF, studies agree on 60 as a good default value
+        # giving good results across many datasets. k allow bigger score
+        # difference among the best results and a smaller score difference among
+        # bad results
+        self._k = k
+        self._keyword_boost = keyword_weight
+        self._semantic_boost = semantic_weight
+
+    @rank_fusion_observer.wrap({"type": "reciprocal_rank_fusion"})
+    def _fuse(
+        self, keyword: Iterable[TextBlockMatch], semantic: Iterable[TextBlockMatch]
+    ) -> list[TextBlockMatch]:
+        scores: dict[ParagraphId, tuple[float, SCORE_TYPE]] = {}
+        match_positions: dict[ParagraphId, list[tuple[int, int]]] = {}
+
+        # sort results by it's score before merging them
+        keyword = [k for k in sorted(keyword, key=lambda r: r.score, reverse=True)]
+        semantic = [s for s in sorted(semantic, key=lambda r: r.score, reverse=True)]
+
+        rankings = [
+            (keyword, self._keyword_boost),
+            (semantic, self._semantic_boost),
+        ]
+        for r, (ranking, boost) in enumerate(rankings):
+            for i, result in enumerate(ranking):
+                id = result.paragraph_id
+                score, score_type = scores.setdefault(id, (0, result.score_type))
+                score += 1 / (self._k + i) * boost
+                if {score_type, result.score_type} == {SCORE_TYPE.BM25, SCORE_TYPE.VECTOR}:
+                    score_type = SCORE_TYPE.BOTH
+                scores[id] = (score, score_type)
+
+                position = (r, i)
+                match_positions.setdefault(result.paragraph_id, []).append(position)
+
+        merged = []
+        for paragraph_id, positions in match_positions.items():
+            # we are getting only one position, effectively deduplicating
+            # multiple matches for the same text block
+            r, i = match_positions[paragraph_id][0]
+            score, score_type = scores[paragraph_id]
+            result = rankings[r][0][i]
+            result.score = score
+            result.score_type = score_type
+            merged.append(result)
+
+        merged.sort(key=lambda x: x.score, reverse=True)
+        return merged
+
+
+def get_rank_fusion(rank_fusion: parser_models.RankFusion) -> RankFusionAlgorithm:
+    """Given a rank fusion API type, return the appropiate rank fusion algorithm instance"""
+    algorithm: RankFusionAlgorithm
+    window = rank_fusion.window
+
+    if isinstance(rank_fusion, parser_models.ReciprocalRankFusion):
+        algorithm = ReciprocalRankFusion(
+            k=rank_fusion.k,
+            window=window,
+            keyword_weight=rank_fusion.boosting.keyword,
+            semantic_weight=rank_fusion.boosting.semantic,
+        )
+
+    else:
+        logger.error(f"Unknown rank fusion algorithm {type(rank_fusion)}: {rank_fusion}. Using default")
+        algorithm = ReciprocalRankFusion(window=window)
+
+    return algorithm