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

nucliadb/search/search/rerankers.py

@@ -0,0 +1,270 @@
+# 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, abstractproperty
+from dataclasses import dataclass
+from typing import Optional
+
+from nucliadb.search.predict import ProxiedPredictAPIError, SendToPredictError
+from nucliadb.search.search.query_parser import models as parser_models
+from nucliadb.search.utilities import get_predict
+from nucliadb_models.internal.predict import RerankModel
+from nucliadb_models.search import (
+    SCORE_TYPE,
+    KnowledgeboxFindResults,
+)
+from nucliadb_telemetry.metrics import Observer
+
+logger = logging.getLogger(__name__)
+
+reranker_observer = Observer("reranker", labels={"type": ""})
+
+
+@dataclass
+class RerankableItem:
+    id: str
+    score: float
+    score_type: SCORE_TYPE
+    content: str
+
+
+@dataclass
+class RankedItem:
+    id: str
+    score: float
+    score_type: SCORE_TYPE
+
+
+@dataclass
+class RerankingOptions:
+    kbid: str
+
+    # Query used to retrieve the results to be reranked. Smart rerankers will use it
+    query: str
+
+
+class Reranker(ABC):
+    @abstractproperty
+    def window(self) -> Optional[int]:
+        """Number of elements the reranker requests. `None` means no specific
+        window is enforced."""
+        ...
+
+    @property
+    def needs_extra_results(self) -> bool:
+        return self.window is not None
+
+    async def rerank(self, items: list[RerankableItem], options: RerankingOptions) -> list[RankedItem]:
+        """Given a query and a set of resources, rerank elements and return the
+        list of reranked items sorted by decreasing score. The list will contain
+        at most, `window` elements.
+
+        """
+        # Enforce reranker window and drop the rest
+        # XXX: other search engines allow a mix of reranked and not reranked
+        # results, there's no technical reason we can't do it
+        items = items[: self.window]
+        reranked = await self._rerank(items, options)
+        return reranked
+
+    @abstractmethod
+    async def _rerank(
+        self, items: list[RerankableItem], options: RerankingOptions
+    ) -> list[RankedItem]: ...
+
+
+class NoopReranker(Reranker):
+    """No-operation reranker. Given a list of items to rerank, it does nothing
+    with them and return the items in the same order. It can be use to not alter
+    the previous ordering.
+
+    """
+
+    @property
+    def window(self) -> Optional[int]:
+        return None
+
+    @reranker_observer.wrap({"type": "noop"})
+    async def _rerank(self, items: list[RerankableItem], options: RerankingOptions) -> list[RankedItem]:
+        return [
+            RankedItem(
+                id=item.id,
+                score=item.score,
+                score_type=item.score_type,
+            )
+            for item in items
+        ]
+
+
+class PredictReranker(Reranker):
+    """Rerank using a reranking model.
+
+    It uses Predict API to rerank elements using a model trained for this
+
+    """
+
+    def __init__(self, window: int):
+        self._window = window
+
+    @property
+    def window(self) -> int:
+        return self._window
+
+    @reranker_observer.wrap({"type": "predict"})
+    async def _rerank(self, items: list[RerankableItem], options: RerankingOptions) -> list[RankedItem]:
+        if len(items) == 0:
+            return []
+
+        predict = get_predict()
+
+        # Conversion to format expected by predict. At the same time,
+        # deduplicates paragraphs found in different indices
+        context = {item.id: item.content for item in items}
+        request = RerankModel(
+            question=options.query,
+            user_id="",  # TODO
+            context=context,
+        )
+        try:
+            response = await predict.rerank(options.kbid, request)
+        except (SendToPredictError, ProxiedPredictAPIError):
+            # predict failed, we can't rerank
+            reranked = [
+                RankedItem(
+                    id=item.id,
+                    score=item.score,
+                    score_type=item.score_type,
+                )
+                for item in items
+            ]
+        else:
+            reranked = [
+                RankedItem(
+                    id=id,
+                    score=score,
+                    score_type=SCORE_TYPE.RERANKER,
+                )
+                for id, score in response.context_scores.items()
+            ]
+        sort_by_score(reranked)
+        best = reranked
+        return best
+
+
+class MultiMatchBoosterReranker(Reranker):
+    """This reranker gives more value to items that come from different indices"""
+
+    @property
+    def window(self) -> Optional[int]:
+        return None
+
+    @reranker_observer.wrap({"type": "multi_match_booster"})
+    async def _rerank(self, items: list[RerankableItem], options: RerankingOptions) -> list[RankedItem]:
+        """Given a list of rerankable items, boost matches that appear multiple
+        times. The returned list can be smaller than the initial, as repeated
+        matches are deduplicated.
+        """
+        reranked_by_id = {}
+        for item in items:
+            if item.id not in reranked_by_id:
+                reranked_by_id[item.id] = RankedItem(
+                    id=item.id,
+                    score=item.score,
+                    score_type=item.score_type,
+                )
+            else:
+                # it's a mutiple match, boost the score
+                if reranked_by_id[item.id].score < item.score:
+                    # previous implementation noted that we are using vector
+                    # score x2 when we find a multiple match. However, this may
+                    # not be true, as the same paragraph could come in any
+                    # position in the rank fusioned result list
+                    reranked_by_id[item.id].score = item.score * 2
+
+                reranked_by_id[item.id].score_type = SCORE_TYPE.BOTH
+
+        reranked = list(reranked_by_id.values())
+        sort_by_score(reranked)
+        return reranked
+
+
+def get_reranker(reranker: parser_models.Reranker) -> Reranker:
+    algorithm: Reranker
+
+    if isinstance(reranker, parser_models.NoopReranker):
+        algorithm = NoopReranker()
+
+    elif isinstance(reranker, parser_models.MultiMatchBoosterReranker):
+        algorithm = MultiMatchBoosterReranker()
+
+    elif isinstance(reranker, parser_models.PredictReranker):
+        algorithm = PredictReranker(reranker.window)
+
+    else:
+        logger.warning(f"Unknown reranker requested: {reranker}. Using default instead")
+        algorithm = MultiMatchBoosterReranker()
+
+    return algorithm
+
+
+def sort_by_score(items: list[RankedItem]):
+    """Sort `items` in place by decreasing score"""
+    items.sort(key=lambda item: item.score, reverse=True)
+
+
+def apply_reranking(results: KnowledgeboxFindResults, reranked: list[RankedItem]):
+    """Given a list of reranked items, update the find results payload.
+
+    *ATENTION* we assume `reranked` is an ordered list of decreasing relevance
+    and contains *only* the items relevant for this response. Any paragraph not
+    found in `reranked` will be removed from the `results`
+
+    """
+    inverted_results = {}
+    for rid, resource in results.resources.items():
+        for field_id, field in resource.fields.items():
+            for paragraph_id, paragraph in field.paragraphs.items():
+                inverted_results[paragraph_id] = (
+                    paragraph,
+                    (field_id, field),
+                    (rid, resource),
+                )
+
+    # update results and best matches according to new scores
+    results.best_matches.clear()
+    for order, item in enumerate(reranked):
+        paragraph_id = item.id
+        paragraph = inverted_results[paragraph_id][0]
+        paragraph.score = item.score
+        paragraph.score_type = item.score_type
+        paragraph.order = order
+        results.best_matches.append(paragraph_id)
+
+    # prune uneeded results (not appearing in `reranked`)
+    extra = set(inverted_results.keys()) - set(results.best_matches)
+    for paragraph_id in extra:
+        _, (field_id, field), (rid, resource) = inverted_results[paragraph_id]
+        field.paragraphs.pop(paragraph_id)
+        if len(field.paragraphs) == 0:
+            resource.fields.pop(field_id)
+
+        if len(resource.fields) == 0:
+            results.resources.pop(rid)

nucliadb/search/search/shards.py

@@ -19,20 +19,15 @@
 #
 import asyncio
 
+from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb_protos.nodereader_pb2 import (
     GetShardRequest,
-    ParagraphSearchRequest,
-    ParagraphSearchResponse,
-    RelationSearchRequest,
-    RelationSearchResponse,
     SearchRequest,
     SearchResponse,
     SuggestRequest,
     SuggestResponse,
 )
 from nucliadb_protos.noderesources_pb2 import Shard
-
-from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb_telemetry import metrics
 
 node_observer = metrics.Observer(
@@ -44,9 +39,7 @@ node_observer = metrics.Observer(
 )
 
 
-async def query_shard(
-    node: AbstractIndexNode, shard: str, query: SearchRequest
-) -> SearchResponse:
+async def query_shard(node: AbstractIndexNode, shard: str, query: SearchRequest) -> SearchResponse:
     req = SearchRequest()
     req.CopyFrom(query)
     req.shard = shard
@@ -61,31 +54,9 @@ async def get_shard(node: AbstractIndexNode, shard_id: str) -> Shard:
         return await node.reader.GetShard(req)  # type: ignore
 
 
-async def query_paragraph_shard(
-    node: AbstractIndexNode, shard: str, query: ParagraphSearchRequest
-) -> ParagraphSearchResponse:
-    req = ParagraphSearchRequest()
-    req.CopyFrom(query)
-    req.id = shard
-    with node_observer({"type": "paragraph_search", "node_id": node.id}):
-        return await node.reader.ParagraphSearch(req)  # type: ignore
-
-
-async def suggest_shard(
-    node: AbstractIndexNode, shard: str, query: SuggestRequest
-) -> SuggestResponse:
+async def suggest_shard(node: AbstractIndexNode, shard: str, query: SuggestRequest) -> SuggestResponse:
     req = SuggestRequest()
     req.CopyFrom(query)
     req.shard = shard
     with node_observer({"type": "suggest", "node_id": node.id}):
         return await node.reader.Suggest(req)  # type: ignore
-
-
-async def relations_shard(
-    node: AbstractIndexNode, shard: str, query: RelationSearchRequest
-) -> RelationSearchResponse:
-    req = RelationSearchRequest()
-    req.CopyFrom(query)
-    req.shard_id = shard
-    with node_observer({"type": "relation_search", "node_id": node.id}):
-        return await node.reader.RelationSearch(req)  # type: ignore

nucliadb/search/search/summarize.py

@@ -20,8 +20,6 @@
 import asyncio
 from typing import Optional
 
-from nucliadb_protos.utils_pb2 import ExtractedText
-
 from nucliadb.common import datamanagers
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.fields.base import Field
@@ -35,6 +33,7 @@ from nucliadb_models.search import (
     SummarizeRequest,
     SummarizeResourceModel,
 )
+from nucliadb_protos.utils_pb2 import ExtractedText
 from nucliadb_utils.utilities import get_storage
 
 ExtractedTexts = list[tuple[str, str, Optional[ExtractedText]]]
@@ -52,15 +51,11 @@ async def summarize(kbid: str, request: SummarizeRequest) -> SummarizedResponse:
     predict_request.user_prompt = request.user_prompt
     predict_request.summary_kind = request.summary_kind
 
-    for uuid_or_slug, field_id, extracted_text in await get_extracted_texts(
-        kbid, request.resources
-    ):
+    for uuid_or_slug, field_id, extracted_text in await get_extracted_texts(kbid, request.resources):
         if extracted_text is None:
             continue
 
-        fields = predict_request.resources.setdefault(
-            uuid_or_slug, SummarizeResourceModel()
-        ).fields
+        fields = predict_request.resources.setdefault(uuid_or_slug, SummarizeResourceModel()).fields
         fields[field_id] = extracted_text.text
 
     if len(predict_request.resources) == 0:
@@ -70,9 +65,7 @@ async def summarize(kbid: str, request: SummarizeRequest) -> SummarizedResponse:
     return await predict.summarize(kbid, predict_request)
 
 
-async def get_extracted_texts(
-    kbid: str, resource_uuids_or_slugs: list[str]
-) -> ExtractedTexts:
+async def get_extracted_texts(kbid: str, resource_uuids_or_slugs: list[str]) -> ExtractedTexts:
     results: ExtractedTexts = []
 
     driver = get_driver()
@@ -82,7 +75,7 @@ async def get_extracted_texts(
     tasks = []
 
     # Schedule getting extracted text for each field of each resource
-    async with driver.transaction() as txn:
+    async with driver.transaction(read_only=True) as txn:
         if not await datamanagers.kb.exists_kb(txn, kbid=kbid):
             raise datamanagers.exceptions.KnowledgeBoxNotFound(kbid)
 
@@ -90,16 +83,12 @@ async def get_extracted_texts(
         for uuid_or_slug in set(resource_uuids_or_slugs):
             uuid = await get_resource_uuid(kb_orm, uuid_or_slug)
             if uuid is None:
-                logger.warning(
-                    f"Resource {uuid_or_slug} not found in KB", extra={"kbid": kbid}
-                )
+                logger.warning(f"Resource {uuid_or_slug} not found in KB", extra={"kbid": kbid})
                 continue
             resource_orm = Resource(txn=txn, storage=storage, kb=kb_orm, uuid=uuid)
             fields = await resource_orm.get_fields(force=True)
             for _, field in fields.items():
-                task = asyncio.create_task(
-                    get_extracted_text(uuid_or_slug, field, max_tasks)
-                )
+                task = asyncio.create_task(get_extracted_text(uuid_or_slug, field, max_tasks))
                 tasks.append(task)
 
         if len(tasks) == 0:

nucliadb/search/search/utils.py

@@ -17,9 +17,26 @@
 # 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 typing import Optional, Union
 
+from pydantic import BaseModel
+
+from nucliadb.common.datamanagers.atomic import kb
 from nucliadb_models.search import BaseSearchRequest, MinScore
+from nucliadb_utils import const
+from nucliadb_utils.utilities import has_feature
+
+logger = logging.getLogger(__name__)
+
+
+async def filter_hidden_resources(kbid: str, show_hidden: bool) -> Optional[bool]:
+    kb_config = await kb.get_config(kbid=kbid)
+    hidden_enabled = kb_config and kb_config.hidden_resources_enabled
+    if hidden_enabled and not show_hidden:
+        return False
+    else:
+        return None  # None = No filtering, show all resources
 
 
 def is_empty_query(request: BaseSearchRequest) -> bool:
@@ -36,7 +53,7 @@ def is_exact_match_only_query(request: BaseSearchRequest) -> bool:
     'foo "something" else' -> False
     """
     query = request.query.strip()
-    return len(query) > 0 and query[0] == '"' and query[-1] == '"'
+    return len(query) > 0 and query.startswith('"') and query.endswith('"')
 
 
 def should_disable_vector_search(request: BaseSearchRequest) -> bool:
@@ -58,9 +75,7 @@ def min_score_from_query_params(
     deprecated_min_score: Optional[float],
 ) -> MinScore:
     # Keep backward compatibility with the deprecated min_score parameter
-    semantic = (
-        deprecated_min_score if min_score_semantic is None else min_score_semantic
-    )
+    semantic = deprecated_min_score if min_score_semantic is None else min_score_semantic
     return MinScore(bm25=min_score_bm25, semantic=semantic)
 
 
@@ -72,3 +87,11 @@ def min_score_from_payload(min_score: Optional[Union[float, MinScore]]) -> MinSc
     elif isinstance(min_score, float):
         return MinScore(bm25=0, semantic=min_score)
     return min_score
+
+
+def maybe_log_request_payload(kbid: str, endpoint: str, item: BaseModel):
+    if has_feature(const.Features.LOG_REQUEST_PAYLOADS, context={"kbid": kbid}, default=False):
+        logger.info(
+            "Request payload",
+            extra={"kbid": kbid, "endpoint": endpoint, "payload": item.model_dump_json()},
+        )

nucliadb/search/settings.py

@@ -18,6 +18,8 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
+from typing import Optional
+
 from pydantic import Field
 
 from nucliadb.ingest.settings import DriverSettings
@@ -28,8 +30,20 @@ class Settings(DriverSettings):
     slow_find_log_threshold: float = Field(
         default=3.0,
         title="Slow query log threshold",
-        description="The threshold in seconds for logging slow queries",
+        description="The threshold in seconds for logging slow find queries",
+    )
+
+    slow_node_query_log_threshold: float = Field(
+        default=2.0,
+        title="Slow node query log threshold",
+        description="The threshold in seconds for logging slow node queries",
+    )
+    prequeries_max_parallel: int = Field(
+        default=2,
+        title="Prequeries max parallel",
+        description="The maximum number of prequeries to run in parallel per /ask request",
     )
+    nidx_address: Optional[str] = Field(default=None)
 
 
 settings = Settings()

nucliadb/standalone/api_router.py

@@ -98,9 +98,7 @@ def get_temp_access_token(request: Request):
         logger.warning(
             "Dynamically generating JWK key. Please set JWK_KEY env variable to avoid this message."
         )
-        settings.jwk_key = orjson.dumps(
-            jwk.JWK.generate(kty="oct", size=256, kid="dyn")
-        ).decode("utf-8")
+        settings.jwk_key = orjson.dumps(jwk.JWK.generate(kty="oct", size=256, kid="dyn")).decode("utf-8")
     jwetoken.add_recipient(jwk.JWK(**orjson.loads(settings.jwk_key)))
     token = jwetoken.serialize(compact=True)
     return JSONResponse({"token": token})
@@ -154,16 +152,14 @@ def introspect_endpoint(request: Request) -> StreamingResponse:
     return StreamingResponse(
         content=introspect.stream_tar(request.app),
         status_code=200,
-        headers={
-            "Content-Disposition": f"attachment; filename=introspect_{introspect_id}.tar.gz"
-        },
+        headers={"Content-Disposition": f"attachment; filename=introspect_{introspect_id}.tar.gz"},
         media_type="application/octet-stream",
     )
 
 
 @standalone_api_router.get("/pull/position")
 async def pull_status(request: Request) -> JSONResponse:
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         # standalone assumes 1 partition
         current_offset = await datamanagers.processing.get_pull_offset(
             txn, pull_type_id=processing.get_nua_api_id(), partition="1"
@@ -180,9 +176,7 @@ class UpdatePullPosition(pydantic.BaseModel):
 
 
 @standalone_api_router.patch("/pull/position")
-async def update_pull_position(
-    request: Request, item: UpdatePullPosition
-) -> JSONResponse:
+async def update_pull_position(request: Request, item: UpdatePullPosition) -> JSONResponse:
     async with datamanagers.with_transaction() as txn:
         # standalone assumes 1 partition
         await datamanagers.processing.set_pull_offset(

nucliadb/standalone/app.py

@@ -20,7 +20,6 @@
 import logging
 import os
 
-import nucliadb_admin_assets  # type: ignore
 from fastapi import FastAPI
 from fastapi.responses import RedirectResponse
 from fastapi.staticfiles import StaticFiles
@@ -31,13 +30,12 @@ from starlette.requests import ClientDisconnect
 from starlette.responses import HTMLResponse
 from starlette.routing import Mount
 
-from nucliadb.common.context.fastapi import set_app_context
+import nucliadb_admin_assets  # type: ignore
 from nucliadb.middleware import ProcessTimeHeaderMiddleware
-from nucliadb.middleware.transaction import ReadOnlyTransactionMiddleware
 from nucliadb.reader import API_PREFIX
 from nucliadb.reader.api.v1.router import api as api_reader_v1
 from nucliadb.search.api.v1.router import api as api_search_v1
-from nucliadb.standalone.lifecycle import finalize, initialize
+from nucliadb.standalone.lifecycle import lifespan
 from nucliadb.train.api.v1.router import api as api_train_v1
 from nucliadb.writer.api.v1.router import api as api_writer_v1
 from nucliadb_telemetry.fastapi import metrics_endpoint
@@ -45,9 +43,11 @@ from nucliadb_telemetry.fastapi.utils import (
     client_disconnect_handler,
     global_exception_handler,
 )
+from nucliadb_utils.audit.stream import AuditMiddleware
 from nucliadb_utils.fastapi.openapi import extend_openapi
 from nucliadb_utils.fastapi.versioning import VersionedFastAPI
 from nucliadb_utils.settings import http_settings, running_settings
+from nucliadb_utils.utilities import get_audit
 
 from .api_router import standalone_api_router
 from .auth import get_auth_backend
@@ -71,7 +71,7 @@ HOMEPAGE_HTML = """
     <h2>Quick Links</h2>
     <ul>
         <li><a href="/admin">Admin UI</a></li>
-        <li><a href="https://docs.nuclia.dev/docs/guides/nucliadb/deploy/basics">NucliaDB Deployment Documentation</a></li>
+        <li><a href="https://docs.nuclia.dev/docs/management/nucliadb/deploy/basics">NucliaDB Deployment Documentation</a></li>
         <li><a href="https://docs.nuclia.dev/docs/api">API Reference</a></li>
         <li><a href="/api/v1/docs">API Explorer</a></li>
         <li><a href="/metrics">Metrics</a></li>
@@ -94,7 +94,7 @@ def application_factory(settings: Settings) -> FastAPI:
             AuthenticationMiddleware,
             backend=get_auth_backend(settings),
         ),
-        Middleware(ReadOnlyTransactionMiddleware),
+        Middleware(AuditMiddleware, audit_utility_getter=get_audit),
     ]
     if running_settings.debug:
         middleware.append(Middleware(ProcessTimeHeaderMiddleware))
@@ -102,8 +102,7 @@ def application_factory(settings: Settings) -> FastAPI:
     fastapi_settings = dict(
         debug=running_settings.debug,
         middleware=middleware,
-        on_startup=[initialize],
-        on_shutdown=[finalize],
+        lifespan=lifespan,
         exception_handlers={
             Exception: global_exception_handler,
             ClientDisconnect: client_disconnect_handler,
@@ -140,9 +139,7 @@ def application_factory(settings: Settings) -> FastAPI:
     # mount admin app assets
     application.mount(
         "/admin",
-        StaticFiles(
-            directory=os.path.dirname(nucliadb_admin_assets.__file__), html=True
-        ),
+        StaticFiles(directory=os.path.dirname(nucliadb_admin_assets.__file__), html=True),
         name="static",
     )
     # redirect /contributor -> /admin
@@ -158,7 +155,4 @@ def application_factory(settings: Settings) -> FastAPI:
         if isinstance(route, Mount):
             route.app.settings = settings  # type: ignore
 
-    # Inject application context into the fastapi app's state
-    set_app_context(application)
-
     return application