nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/search/api/v1/retrieve.py

@@ -0,0 +1,130 @@
+# 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 time import time
+
+from fastapi import Header, HTTPException, Request
+from fastapi_versioning import version
+
+from nucliadb.common.exceptions import InvalidQueryError
+from nucliadb.common.external_index_providers.base import TextBlockMatch
+from nucliadb.common.models_utils import to_proto
+from nucliadb.search.api.v1.router import KB_PREFIX, api
+from nucliadb.search.search.query_parser.parsers.retrieve import parse_retrieve
+from nucliadb.search.search.retrieval import text_block_search
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.retrieval import (
+    Metadata,
+    RetrievalMatch,
+    RetrievalRequest,
+    RetrievalResponse,
+    Scores,
+)
+from nucliadb_models.search import NucliaDBClientType
+from nucliadb_utils.authentication import requires
+from nucliadb_utils.utilities import get_audit
+
+
+@api.post(
+    f"/{KB_PREFIX}/{{kbid}}/retrieve",
+    status_code=200,
+    description="Search text blocks on a Knowledge Box",
+    include_in_schema=False,
+    tags=["Search"],
+)
+@requires(NucliaDBRoles.READER)
+@version(1)
+async def _retrieve_endpoint(
+    request: Request,
+    kbid: str,
+    item: RetrievalRequest,
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_nucliadb_user: str = Header(""),
+    x_forwarded_for: str = Header(""),
+) -> RetrievalResponse:
+    return await retrieve_endpoint(
+        kbid,
+        item,
+        x_ndb_client=x_ndb_client,
+        x_nucliadb_user=x_nucliadb_user,
+        x_forwarded_for=x_forwarded_for,
+    )
+
+
+async def retrieve_endpoint(
+    kbid: str,
+    item: RetrievalRequest,
+    *,
+    x_ndb_client: NucliaDBClientType,
+    x_nucliadb_user: str,
+    x_forwarded_for: str,
+) -> RetrievalResponse:
+    audit = get_audit()
+    start_time = time()
+
+    try:
+        retrieval = await parse_retrieve(kbid, item)
+    except InvalidQueryError as err:
+        raise HTTPException(
+            status_code=422,
+            detail=str(err),
+        )
+
+    text_blocks, pb_query, _, _ = await text_block_search(kbid, retrieval)
+
+    # cut the top K, we may have more due to extra results used for rank fusion
+    text_blocks = text_blocks[: retrieval.top_k]
+
+    # convert to response models
+    matches = [text_block_match_to_retrieval_match(text_block) for text_block in text_blocks]
+
+    if audit is not None:
+        retrieval_time = time() - start_time
+        audit.retrieve(
+            kbid,
+            x_nucliadb_user,
+            to_proto.client_type(x_ndb_client),
+            x_forwarded_for,
+            retrieval_time,
+            # TODO(decoupled-ask): add interesting things to audit
+        )
+
+    return RetrievalResponse(matches=matches)
+
+
+def text_block_match_to_retrieval_match(item: TextBlockMatch) -> RetrievalMatch:
+    return RetrievalMatch(
+        id=item.paragraph_id.full(),
+        score=Scores(
+            value=item.current_score.score,
+            source=item.current_score.source,
+            type=item.current_score.type,
+            history=item.scores,
+        ),
+        metadata=Metadata(
+            field_labels=item.field_labels,
+            paragraph_labels=item.paragraph_labels,
+            is_an_image=item.is_an_image,
+            is_a_table=item.is_a_table,
+            source_file=item.representation_file,
+            page=item.position.page_number,
+            in_page_with_visual=item.page_with_visual,
+        ),
+    )

nucliadb/search/api/v1/search.py

@@ -19,7 +19,6 @@
 #
 import json
 from time import time
-from typing import Optional, Union
 
 from fastapi import Body, Header, Query, Request, Response
 from fastapi.openapi.models import Example
@@ -37,7 +36,10 @@ from nucliadb.search.requesters.utils import Method, nidx_query
 from nucliadb.search.search import cache
 from nucliadb.search.search.merge import merge_results
 from nucliadb.search.search.query_parser.parsers.search import parse_search
-from nucliadb.search.search.query_parser.parsers.unit_retrieval import legacy_convert_retrieval_to_proto
+from nucliadb.search.search.query_parser.parsers.unit_retrieval import (
+    convert_retrieval_to_proto,
+    is_incomplete,
+)
 from nucliadb.search.search.utils import (
     min_score_from_query_params,
 )
@@ -65,7 +67,7 @@ from nucliadb_utils.utilities import get_audit
 SEARCH_EXAMPLES = {
     "filtering_by_icon": Example(
         summary="Search for pdf documents where the text 'Noam Chomsky' appears",
-        description="For a complete list of filters, visit: https://github.com/nuclia/nucliadb/blob/main/docs/internal/SEARCH.md#filters-and-facets",  # noqa
+        description="For a complete list of filters, visit: https://github.com/nuclia/nucliadb/blob/main/docs/internal/SEARCH.md#filters-and-facets",
         value={
             "query": "Noam Chomsky",
             "filters": ["/icon/application/pdf"],
@@ -74,7 +76,7 @@ SEARCH_EXAMPLES = {
     ),
     "get_language_counts": Example(
         summary="Get the number of documents for each language",
-        description="For a complete list of facets, visit: https://github.com/nuclia/nucliadb/blob/main/docs/internal/SEARCH.md#filters-and-facets",  # noqa
+        description="For a complete list of facets, visit: https://github.com/nuclia/nucliadb/blob/main/docs/internal/SEARCH.md#filters-and-facets",
         value={
             "page_size": 0,
             "faceted": ["/s/p"],
@@ -88,7 +90,7 @@ SEARCH_EXAMPLES = {
     f"/{KB_PREFIX}/{{kbid}}/search",
     status_code=200,
     summary="Search Knowledge Box",
-    description="Search on a Knowledge Box and retrieve separate results for documents, paragraphs, and sentences. Usually, it is better to use `find`",  # noqa: E501
+    description="Search on a Knowledge Box and retrieve separate results for documents, paragraphs, and sentences. Usually, it is better to use `find`",
     response_model=KnowledgeboxSearchResults,
     response_model_exclude_unset=True,
     tags=["Search"],
@@ -100,37 +102,35 @@ async def search_knowledgebox(
     response: Response,
     kbid: str,
     query: str = fastapi_query(SearchParamDefaults.query),
-    filter_expression: Optional[str] = fastapi_query(SearchParamDefaults.filter_expression),
+    filter_expression: str | None = fastapi_query(SearchParamDefaults.filter_expression),
     fields: list[str] = fastapi_query(SearchParamDefaults.fields),
     filters: list[str] = fastapi_query(SearchParamDefaults.filters),
     faceted: list[str] = fastapi_query(SearchParamDefaults.faceted),
     sort_field: SortField = fastapi_query(SearchParamDefaults.sort_field),
-    sort_limit: Optional[int] = fastapi_query(SearchParamDefaults.sort_limit),
     sort_order: SortOrder = fastapi_query(SearchParamDefaults.sort_order),
     top_k: int = fastapi_query(SearchParamDefaults.top_k),
-    min_score: Optional[float] = Query(
+    offset: int = fastapi_query(SearchParamDefaults.offset),
+    min_score: float | None = Query(
         default=None,
-        description="Minimum similarity score to filter vector index results. If not specified, the default minimum score of the semantic model associated to the Knowledge Box will be used. Check out the documentation for more information on how to use this parameter: https://docs.nuclia.dev/docs/rag/advanced/search#minimum-score",  # noqa: E501
+        description="Minimum similarity score to filter vector index results. If not specified, the default minimum score of the semantic model associated to the Knowledge Box will be used. Check out the documentation for more information on how to use this parameter: https://docs.nuclia.dev/docs/rag/advanced/search#minimum-score",
         deprecated=True,
     ),
-    min_score_semantic: Optional[float] = Query(
+    min_score_semantic: float | None = Query(
         default=None,
-        description="Minimum semantic similarity score to filter vector index results. If not specified, the default minimum score of the semantic model associated to the Knowledge Box will be used. Check out the documentation for more information on how to use this parameter: https://docs.nuclia.dev/docs/rag/advanced/search#minimum-score",  # noqa: E501
+        description="Minimum semantic similarity score to filter vector index results. If not specified, the default minimum score of the semantic model associated to the Knowledge Box will be used. Check out the documentation for more information on how to use this parameter: https://docs.nuclia.dev/docs/rag/advanced/search#minimum-score",
     ),
     min_score_bm25: float = Query(
         default=0,
         description="Minimum bm25 score to filter paragraph and document index results",
         ge=0,
     ),
-    vectorset: Optional[str] = fastapi_query(SearchParamDefaults.vectorset),
-    range_creation_start: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_start),
-    range_creation_end: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_end),
-    range_modification_start: Optional[DateTime] = fastapi_query(
+    vectorset: str | None = fastapi_query(SearchParamDefaults.vectorset),
+    range_creation_start: DateTime | None = fastapi_query(SearchParamDefaults.range_creation_start),
+    range_creation_end: DateTime | None = fastapi_query(SearchParamDefaults.range_creation_end),
+    range_modification_start: DateTime | None = fastapi_query(
         SearchParamDefaults.range_modification_start
     ),
-    range_modification_end: Optional[DateTime] = fastapi_query(
-        SearchParamDefaults.range_modification_end
-    ),
+    range_modification_end: DateTime | None = fastapi_query(SearchParamDefaults.range_modification_end),
     features: list[SearchOptions] = fastapi_query(
         SearchParamDefaults.search_features,
         default=[
@@ -148,13 +148,12 @@ async def search_knowledgebox(
     extracted: list[ExtractedDataTypeName] = fastapi_query(SearchParamDefaults.extracted),
     with_duplicates: bool = fastapi_query(SearchParamDefaults.with_duplicates),
     with_synonyms: bool = fastapi_query(SearchParamDefaults.with_synonyms),
-    autofilter: bool = fastapi_query(SearchParamDefaults.autofilter),
     security_groups: list[str] = fastapi_query(SearchParamDefaults.security_groups),
     show_hidden: bool = fastapi_query(SearchParamDefaults.show_hidden),
     x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
     x_nucliadb_user: str = Header(""),
     x_forwarded_for: str = Header(""),
-) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
+) -> KnowledgeboxSearchResults | HTTPClientError:
     try:
         expr = FilterExpression.model_validate_json(filter_expression) if filter_expression else None
 
@@ -167,11 +166,7 @@ async def search_knowledgebox(
             fields=fields,
             filters=filters,
             faceted=faceted,
-            sort=(
-                SortOptions(field=sort_field, limit=sort_limit, order=sort_order)
-                if sort_field is not None
-                else None
-            ),
+            sort=(SortOptions(field=sort_field, order=sort_order) if sort_field is not None else None),
             top_k=top_k,
             min_score=min_score_from_query_params(min_score_bm25, min_score_semantic, min_score),
             vectorset=vectorset,
@@ -187,9 +182,9 @@ async def search_knowledgebox(
             extracted=extracted,
             with_duplicates=with_duplicates,
             with_synonyms=with_synonyms,
-            autofilter=autofilter,
             security=security,
             show_hidden=show_hidden,
+            offset=offset,
         )
     except ValidationError as exc:
         detail = json.loads(exc.json())
@@ -201,7 +196,7 @@ async def search_knowledgebox(
     f"/{KB_PREFIX}/{{kbid}}/search",
     status_code=200,
     summary="Search Knowledge Box",
-    description="Search on a Knowledge Box and retrieve separate results for documents, paragraphs, and sentences. Usually, it is better to use `find`",  # noqa: E501
+    description="Search on a Knowledge Box and retrieve separate results for documents, paragraphs, and sentences. Usually, it is better to use `find`",
     response_model=KnowledgeboxSearchResults,
     response_model_exclude_unset=True,
     tags=["Search"],
@@ -216,7 +211,7 @@ async def search_post_knowledgebox(
     x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
     x_nucliadb_user: str = Header(""),
     x_forwarded_for: str = Header(""),
-) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
+) -> KnowledgeboxSearchResults | HTTPClientError:
     return await _search_endpoint(response, kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for)
 
 
@@ -228,7 +223,7 @@ async def _search_endpoint(
     x_nucliadb_user: str,
     x_forwarded_for: str,
     **kwargs,
-) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
+) -> KnowledgeboxSearchResults | HTTPClientError:
     try:
         with cache.request_caches():
             results, incomplete = await search(
@@ -256,13 +251,14 @@ async def search(
     x_nucliadb_user: str,
     x_forwarded_for: str,
     do_audit: bool = True,
-    with_status: Optional[ResourceProcessingStatus] = None,
+    with_status: ResourceProcessingStatus | None = None,
 ) -> tuple[KnowledgeboxSearchResults, bool]:
     audit = get_audit()
     start_time = time()
 
     parsed = await parse_search(kbid, item)
-    pb_query, incomplete_results, autofilters, _ = await legacy_convert_retrieval_to_proto(parsed)
+    incomplete_results = is_incomplete(parsed.retrieval)
+    pb_query = convert_retrieval_to_proto(parsed.retrieval)
 
     # We need to query all nodes
     results, queried_shards = await nidx_query(kbid, Method.SEARCH, pb_query)
@@ -276,6 +272,7 @@ async def search(
         field_type_filter=item.field_type_filter,
         extracted=item.extracted,
         highlight=item.highlight,
+        offset=item.offset,
     )
 
     if audit is not None and do_audit:
@@ -290,5 +287,4 @@ async def search(
         )
 
     search_results.shards = queried_shards
-    search_results.autofilters = autofilters
     return search_results, incomplete_results

nucliadb/search/api/v1/suggest.py

@@ -19,7 +19,6 @@
 #
 import json
 from datetime import datetime
-from typing import Optional, Union
 
 from fastapi import Header, Request, Response
 from fastapi_versioning import version
@@ -64,20 +63,18 @@ async def suggest_knowledgebox(
     response: Response,
     kbid: str,
     query: str = fastapi_query(SearchParamDefaults.suggest_query),
-    filter_expression: Optional[str] = fastapi_query(
+    filter_expression: str | None = fastapi_query(
         SearchParamDefaults.filter_expression, include_in_schema=False
     ),
     fields: list[str] = fastapi_query(SearchParamDefaults.fields),
     filters: list[str] = fastapi_query(SearchParamDefaults.filters),
     faceted: list[str] = fastapi_query(SearchParamDefaults.faceted),
-    range_creation_start: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_start),
-    range_creation_end: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_end),
-    range_modification_start: Optional[DateTime] = fastapi_query(
+    range_creation_start: DateTime | None = fastapi_query(SearchParamDefaults.range_creation_start),
+    range_creation_end: DateTime | None = fastapi_query(SearchParamDefaults.range_creation_end),
+    range_modification_start: DateTime | None = fastapi_query(
         SearchParamDefaults.range_modification_start
     ),
-    range_modification_end: Optional[DateTime] = fastapi_query(
-        SearchParamDefaults.range_modification_end
-    ),
+    range_modification_end: DateTime | None = fastapi_query(SearchParamDefaults.range_modification_end),
     features: list[SuggestOptions] = fastapi_query(SearchParamDefaults.suggest_features),
     show: list[ResourceProperties] = fastapi_query(SearchParamDefaults.show),
     field_type_filter: list[FieldTypeName] = fastapi_query(
@@ -89,7 +86,7 @@ async def suggest_knowledgebox(
     debug: bool = fastapi_query(SearchParamDefaults.debug),
     highlight: bool = fastapi_query(SearchParamDefaults.highlight),
     show_hidden: bool = fastapi_query(SearchParamDefaults.show_hidden),
-) -> Union[KnowledgeboxSuggestResults, HTTPClientError]:
+) -> KnowledgeboxSuggestResults | HTTPClientError:
     try:
         expr = FilterExpression.model_validate_json(filter_expression) if filter_expression else None
 
@@ -126,14 +123,14 @@ async def suggest(
     response,
     kbid: str,
     query: str,
-    filter_expression: Optional[FilterExpression],
+    filter_expression: FilterExpression | None,
     fields: list[str],
     filters: list[str],
     faceted: list[str],
-    range_creation_start: Optional[datetime],
-    range_creation_end: Optional[datetime],
-    range_modification_start: Optional[datetime],
-    range_modification_end: Optional[datetime],
+    range_creation_start: datetime | None,
+    range_creation_end: datetime | None,
+    range_modification_start: datetime | None,
+    range_modification_end: datetime | None,
     features: list[SuggestOptions],
     show: list[ResourceProperties],
     field_type_filter: list[FieldTypeName],

nucliadb/search/api/v1/summarize.py

@@ -17,7 +17,6 @@
 # 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 Union
 
 from fastapi import Header, Request
 from fastapi_versioning import version
@@ -48,7 +47,7 @@ async def summarize_endpoint(
     kbid: str,
     item: SummarizeRequest,
     x_show_consumption: bool = Header(default=False),
-) -> Union[SummarizedResponse, HTTPClientError]:
+) -> SummarizedResponse | HTTPClientError:
     try:
         return await summarize(
             kbid=kbid,

nucliadb/search/api/v1/utils.py

@@ -17,7 +17,7 @@
 # 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, Optional
+from typing import Any
 
 from fastapi import Query
 
@@ -26,7 +26,7 @@ from nucliadb_models.search import ParamDefault
 _NOT_SET = object()
 
 
-def fastapi_query(param: ParamDefault, default: Optional[Any] = _NOT_SET, **kw) -> Query:  # type: ignore
+def fastapi_query(param: ParamDefault, default: Any | None = _NOT_SET, **kw) -> Query:  # type: ignore
     # Be able to override default value
     if default is _NOT_SET:
         default_value = param.default

nucliadb/search/app.py

@@ -26,7 +26,7 @@ from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.requests import ClientDisconnect, Request
 from starlette.responses import HTMLResponse
 
-from nucliadb.middleware import ProcessTimeHeaderMiddleware
+from nucliadb.middleware import ClientErrorPayloadLoggerMiddleware, ProcessTimeHeaderMiddleware
 from nucliadb.search import API_PREFIX
 from nucliadb.search.api.v1.router import api as api_v1
 from nucliadb.search.lifecycle import lifespan
@@ -47,6 +47,7 @@ middleware.extend(
     [
         Middleware(AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()),
         Middleware(AuditMiddleware, audit_utility_getter=get_audit),
+        Middleware(ClientErrorPayloadLoggerMiddleware),
     ]
 )
 
@@ -58,7 +59,6 @@ errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
 fastapi_settings = dict(
     debug=running_settings.debug,
-    middleware=middleware,
     lifespan=lifespan,
     exception_handlers={
         Exception: global_exception_handler,
@@ -78,6 +78,7 @@ application = VersionedFastAPI(
     prefix_format=f"/{API_PREFIX}/v{{major}}",
     default_version=(1, 0),
     enable_latest=False,
+    middleware=middleware,
     kwargs=fastapi_settings,
 )
 

nucliadb/search/augmentor/__init__.py

@@ -0,0 +1,21 @@
+# 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 . import fields, paragraphs, resources  # noqa: F401
+from .augmentor import augment  # noqa: F401