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

nucliadb/search/api/v1/resource/search.py

@@ -17,8 +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 datetime import datetime
-from typing import Optional, Union
+from typing import Optional, Union, cast
 
 from fastapi import Header, Request, Response
 from fastapi_versioning import version
@@ -27,20 +26,19 @@ from nucliadb.models.responses import HTTPClientError
 from nucliadb.search.api.v1.router import KB_PREFIX, RESOURCE_PREFIX, api
 from nucliadb.search.api.v1.utils import fastapi_query
 from nucliadb.search.requesters.utils import Method, debug_nodes_info, node_query
+from nucliadb.search.search import cache
 from nucliadb.search.search.exceptions import InvalidQueryError
 from nucliadb.search.search.merge import merge_paragraphs_results
 from nucliadb.search.search.query import paragraph_query_to_pb
-from nucliadb_models.common import FieldTypeName
-from nucliadb_models.resource import ExtractedDataTypeName, NucliaDBRoles
+from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.search import (
     NucliaDBClientType,
-    ResourceProperties,
     ResourceSearchResults,
-    SearchOptions,
     SearchParamDefaults,
     SortField,
     SortOrder,
 )
+from nucliadb_models.utils import DateTime
 from nucliadb_utils.authentication import requires_one
 
 
@@ -64,81 +62,62 @@ async def resource_search(
     fields: list[str] = fastapi_query(SearchParamDefaults.fields),
     filters: list[str] = fastapi_query(SearchParamDefaults.filters),
     faceted: list[str] = fastapi_query(SearchParamDefaults.faceted),
-    sort: Optional[SortField] = fastapi_query(
-        SearchParamDefaults.sort_field, alias="sort_field"
-    ),
+    sort: Optional[SortField] = fastapi_query(SearchParamDefaults.sort_field, alias="sort_field"),
     sort_order: SortOrder = fastapi_query(SearchParamDefaults.sort_order),
-    page_number: int = fastapi_query(SearchParamDefaults.page_number),
-    page_size: int = fastapi_query(SearchParamDefaults.page_size),
-    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(
+    top_k: Optional[int] = fastapi_query(SearchParamDefaults.top_k),
+    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(
         SearchParamDefaults.range_modification_start
     ),
-    range_modification_end: Optional[datetime] = fastapi_query(
+    range_modification_end: Optional[DateTime] = fastapi_query(
         SearchParamDefaults.range_modification_end
     ),
     highlight: bool = fastapi_query(SearchParamDefaults.highlight),
-    show: list[ResourceProperties] = fastapi_query(
-        SearchParamDefaults.show, default=list(ResourceProperties)
-    ),
-    field_type_filter: list[FieldTypeName] = fastapi_query(
-        SearchParamDefaults.field_type_filter, alias="field_type"
-    ),
-    extracted: list[ExtractedDataTypeName] = fastapi_query(
-        SearchParamDefaults.extracted
-    ),
     x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
     debug: bool = fastapi_query(SearchParamDefaults.debug),
     shards: list[str] = fastapi_query(SearchParamDefaults.shards),
 ) -> Union[ResourceSearchResults, HTTPClientError]:
-    # We need to query all nodes
-    try:
-        pb_query = await paragraph_query_to_pb(
-            kbid,
-            [SearchOptions.PARAGRAPH],
-            rid,
-            query,
-            fields,
-            filters,
-            faceted,
-            page_number,
-            page_size,
-            range_creation_start,
-            range_creation_end,
-            range_modification_start,
-            range_modification_end,
-            sort=sort.value if sort else None,
-            sort_ord=sort_order.value,
-        )
-    except InvalidQueryError as exc:
-        return HTTPClientError(status_code=412, detail=str(exc))
+    top_k = top_k or SearchParamDefaults.top_k  # type: ignore
+    top_k = cast(int, top_k)
 
-    results, incomplete_results, queried_nodes = await node_query(
-        kbid, Method.PARAGRAPH, pb_query, shards
-    )
+    with cache.request_caches():
+        try:
+            pb_query = await paragraph_query_to_pb(
+                kbid,
+                rid,
+                query,
+                fields,
+                filters,
+                faceted,
+                top_k,
+                range_creation_start,
+                range_creation_end,
+                range_modification_start,
+                range_modification_end,
+                sort=sort.value if sort else None,
+                sort_ord=sort_order.value,
+            )
+        except InvalidQueryError as exc:
+            return HTTPClientError(status_code=412, detail=str(exc))
 
-    # We need to merge
-    search_results = await merge_paragraphs_results(
-        results,
-        count=page_size,
-        page=page_number,
-        kbid=kbid,
-        show=show,
-        field_type_filter=field_type_filter,
-        extracted=extracted,
-        highlight_split=highlight,
-        min_score=0.0,
-    )
+        results, incomplete_results, queried_nodes = await node_query(
+            kbid, Method.SEARCH, pb_query, shards
+        )
+
+        # We need to merge
+        search_results = await merge_paragraphs_results(
+            results,
+            top_k=top_k,
+            kbid=kbid,
+            highlight_split=highlight,
+            min_score=0.0,
+        )
 
-    response.status_code = 206 if incomplete_results else 200
-    if debug:
-        search_results.nodes = debug_nodes_info(queried_nodes)
+        response.status_code = 206 if incomplete_results else 200
+        if debug:
+            search_results.nodes = debug_nodes_info(queried_nodes)
 
-    queried_shards = [shard_id for _, shard_id in queried_nodes]
-    search_results.shards = queried_shards
-    return search_results
+        queried_shards = [shard_id for _, shard_id in queried_nodes]
+        search_results.shards = queried_shards
+        return search_results

nucliadb/search/api/v1/search.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import json
-from datetime import datetime
 from time import time
 from typing import Optional, Union
 
@@ -28,15 +27,18 @@ from fastapi_versioning import version
 from pydantic import ValidationError
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
+from nucliadb.common.models_utils import to_proto
 from nucliadb.models.responses import HTTPClientError
 from nucliadb.search import predict
 from nucliadb.search.api.v1.router import KB_PREFIX, api
 from nucliadb.search.api.v1.utils import fastapi_query
 from nucliadb.search.requesters.utils import Method, debug_nodes_info, node_query
+from nucliadb.search.search import cache
 from nucliadb.search.search.exceptions import InvalidQueryError
 from nucliadb.search.search.merge import merge_results
 from nucliadb.search.search.query import QueryParser
 from nucliadb.search.search.utils import (
+    filter_hidden_resources,
     min_score_from_payload,
     min_score_from_query_params,
     should_disable_vector_search,
@@ -45,9 +47,7 @@ from nucliadb_models.common import FieldTypeName
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_models.resource import ExtractedDataTypeName, NucliaDBRoles
 from nucliadb_models.search import (
-    CatalogRequest,
     KnowledgeboxSearchResults,
-    MinScore,
     NucliaDBClientType,
     ResourceProperties,
     SearchOptions,
@@ -58,6 +58,7 @@ from nucliadb_models.search import (
     SortOrder,
 )
 from nucliadb_models.security import RequestSecurity
+from nucliadb_models.utils import DateTime
 from nucliadb_utils.authentication import requires
 from nucliadb_utils.exceptions import LimitsExceededError
 from nucliadb_utils.utilities import get_audit
@@ -69,7 +70,7 @@ SEARCH_EXAMPLES = {
         value={
             "query": "Noam Chomsky",
             "filters": ["/icon/application/pdf"],
-            "features": [SearchOptions.DOCUMENT],
+            "features": [SearchOptions.FULLTEXT],
         },
     ),
     "get_language_counts": Example(
@@ -78,7 +79,7 @@ SEARCH_EXAMPLES = {
         value={
             "page_size": 0,
             "faceted": ["/s/p"],
-            "features": [SearchOptions.DOCUMENT],
+            "features": [SearchOptions.FULLTEXT],
         },
     ),
 }
@@ -106,40 +107,36 @@ async def search_knowledgebox(
     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),
-    page_number: int = fastapi_query(SearchParamDefaults.page_number),
-    page_size: int = fastapi_query(SearchParamDefaults.page_size),
+    top_k: int = fastapi_query(SearchParamDefaults.top_k),
     min_score: Optional[float] = 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/docs/using/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",  # noqa: E501
         deprecated=True,
     ),
     min_score_semantic: Optional[float] = 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/docs/using/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",  # noqa: E501
     ),
     min_score_bm25: float = Query(
         default=0,
         description="Minimum bm25 score to filter paragraph and document index results",
         ge=0,
     ),
-    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: 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(
         SearchParamDefaults.range_modification_start
     ),
-    range_modification_end: Optional[datetime] = fastapi_query(
+    range_modification_end: Optional[DateTime] = fastapi_query(
         SearchParamDefaults.range_modification_end
     ),
     features: list[SearchOptions] = fastapi_query(
         SearchParamDefaults.search_features,
         default=[
-            SearchOptions.PARAGRAPH,
-            SearchOptions.DOCUMENT,
-            SearchOptions.VECTOR,
+            SearchOptions.KEYWORD,
+            SearchOptions.FULLTEXT,
+            SearchOptions.SEMANTIC,
         ],
     ),
     debug: bool = fastapi_query(SearchParamDefaults.debug),
@@ -148,14 +145,13 @@ async def search_knowledgebox(
     field_type_filter: list[FieldTypeName] = fastapi_query(
         SearchParamDefaults.field_type_filter, alias="field_type"
     ),
-    extracted: list[ExtractedDataTypeName] = fastapi_query(
-        SearchParamDefaults.extracted
-    ),
+    extracted: list[ExtractedDataTypeName] = fastapi_query(SearchParamDefaults.extracted),
     shards: list[str] = fastapi_query(SearchParamDefaults.shards),
     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(""),
@@ -174,11 +170,9 @@ async def search_knowledgebox(
                 if sort_field is not None
                 else None
             ),
-            page_number=page_number,
-            page_size=page_size,
-            min_score=min_score_from_query_params(
-                min_score_bm25, min_score_semantic, min_score
-            ),
+            top_k=top_k,
+            min_score=min_score_from_query_params(min_score_bm25, min_score_semantic, min_score),
+            vectorset=vectorset,
             range_creation_end=range_creation_end,
             range_creation_start=range_creation_start,
             range_modification_end=range_modification_end,
@@ -194,173 +188,12 @@ async def search_knowledgebox(
             with_synonyms=with_synonyms,
             autofilter=autofilter,
             security=security,
+            show_hidden=show_hidden,
         )
     except ValidationError as exc:
         detail = json.loads(exc.json())
         return HTTPClientError(status_code=422, detail=detail)
-    return await _search_endpoint(
-        response, kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for
-    )
-
-
-@api.get(
-    f"/{KB_PREFIX}/{{kbid}}/catalog",
-    status_code=200,
-    summary="List resources of a Knowledge Box",
-    description="List resources of a Knowledge Box",
-    response_model=KnowledgeboxSearchResults,
-    response_model_exclude_unset=True,
-    tags=["Search"],
-)
-@requires(NucliaDBRoles.READER)
-@version(1)
-async def catalog_get(
-    request: Request,
-    response: Response,
-    kbid: str,
-    query: str = fastapi_query(SearchParamDefaults.query),
-    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),
-    page_number: int = fastapi_query(SearchParamDefaults.page_number),
-    page_size: int = fastapi_query(SearchParamDefaults.page_size),
-    shards: list[str] = fastapi_query(SearchParamDefaults.shards),
-    with_status: Optional[ResourceProcessingStatus] = fastapi_query(
-        SearchParamDefaults.with_status
-    ),
-    debug: bool = fastapi_query(SearchParamDefaults.debug),
-    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(
-        SearchParamDefaults.range_modification_start
-    ),
-    range_modification_end: Optional[datetime] = fastapi_query(
-        SearchParamDefaults.range_modification_end
-    ),
-) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
-    item = CatalogRequest(
-        query=query,
-        filters=filters,
-        faceted=faceted,
-        page_number=page_number,
-        page_size=page_size,
-        shards=shards,
-        debug=debug,
-        with_status=with_status,
-        range_creation_start=range_creation_start,
-        range_creation_end=range_creation_end,
-        range_modification_start=range_modification_start,
-        range_modification_end=range_modification_end,
-    )
-    if sort_field:
-        item.sort = SortOptions(field=sort_field, limit=sort_limit, order=sort_order)
-    return await catalog(kbid, item)
-
-
-@api.post(
-    f"/{KB_PREFIX}/{{kbid}}/catalog",
-    status_code=200,
-    summary="List resources of a Knowledge Box",
-    description="List resources of a Knowledge Box",
-    response_model=KnowledgeboxSearchResults,
-    response_model_exclude_unset=True,
-    tags=["Search"],
-)
-@requires(NucliaDBRoles.READER)
-@version(1)
-async def catalog_post(
-    request: Request,
-    kbid: str,
-    item: CatalogRequest,
-) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
-    return await catalog(kbid, item)
-
-
-async def catalog(
-    kbid: str,
-    item: CatalogRequest,
-):
-    """
-    Catalog endpoint is a simplified version of the search endpoint, it only
-    returns bm25 results on titles and it does not support vector search.
-    It is useful for listing resources in a knowledge box.
-    """
-    try:
-        sort = item.sort
-        if item.sort is None:
-            # By default we sort by creation date (most recent first)
-            sort = SortOptions(
-                field=SortField.CREATED,
-                order=SortOrder.DESC,
-                limit=None,
-            )
-
-        query_parser = QueryParser(
-            kbid=kbid,
-            features=[SearchOptions.DOCUMENT],
-            query=item.query,
-            filters=item.filters,
-            faceted=item.faceted,
-            sort=sort,
-            page_number=item.page_number,
-            page_size=item.page_size,
-            min_score=MinScore(bm25=0, semantic=0),
-            fields=["a/title"],
-            with_status=item.with_status,
-            range_creation_start=item.range_creation_start,
-            range_creation_end=item.range_creation_end,
-            range_modification_start=item.range_modification_start,
-            range_modification_end=item.range_modification_end,
-        )
-        pb_query, _, _ = await query_parser.parse()
-
-        (results, _, queried_nodes) = await node_query(
-            kbid,
-            Method.SEARCH,
-            pb_query,
-            target_shard_replicas=item.shards,
-            # Catalog should not go to read replicas because we want it to be
-            # consistent and most up to date results
-            use_read_replica_nodes=False,
-        )
-
-        # We need to merge
-        search_results = await merge_results(
-            results,
-            count=item.page_size,
-            page=item.page_number,
-            kbid=kbid,
-            show=[ResourceProperties.BASIC],
-            field_type_filter=[],
-            extracted=[],
-            sort=sort,
-            requested_relations=pb_query.relation_subgraph,
-            min_score=query_parser.min_score,
-            highlight=False,
-        )
-        # We don't need sentences, paragraphs or relations on the catalog
-        # response, so we set to None so that fastapi doesn't include them
-        # in the response payload
-        search_results.sentences = None
-        search_results.paragraphs = None
-        search_results.relations = None
-        if item.debug:
-            search_results.nodes = debug_nodes_info(queried_nodes)
-        queried_shards = [shard_id for _, shard_id in queried_nodes]
-        search_results.shards = queried_shards
-        return search_results
-    except InvalidQueryError as exc:
-        return HTTPClientError(status_code=412, detail=str(exc))
-    except KnowledgeBoxNotFound:
-        return HTTPClientError(status_code=404, detail="Knowledge Box not found")
-    except LimitsExceededError as exc:
-        return HTTPClientError(status_code=exc.status_code, detail=exc.detail)
+    return await _search_endpoint(response, kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for)
 
 
 @api.post(
@@ -383,9 +216,7 @@ async def search_post_knowledgebox(
     x_nucliadb_user: str = Header(""),
     x_forwarded_for: str = Header(""),
 ) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
-    return await _search_endpoint(
-        response, kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for
-    )
+    return await _search_endpoint(response, kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for)
 
 
 async def _search_endpoint(
@@ -397,13 +228,13 @@ async def _search_endpoint(
     x_forwarded_for: str,
     **kwargs,
 ) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
-    # All endpoint logic should be here
     try:
-        results, incomplete = await search(
-            kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for, **kwargs
-        )
-        response.status_code = 206 if incomplete else 200
-        return results
+        with cache.request_caches():
+            results, incomplete = await search(
+                kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for, **kwargs
+            )
+            response.status_code = 206 if incomplete else 200
+            return results
     except KnowledgeBoxNotFound:
         return HTTPClientError(status_code=404, detail="Knowledge Box not found")
     except LimitsExceededError as exc:
@@ -431,20 +262,20 @@ async def search(
 
     item.min_score = min_score_from_payload(item.min_score)
 
-    if SearchOptions.VECTOR in item.features:
+    if SearchOptions.SEMANTIC in item.features:
         if should_disable_vector_search(item):
-            item.features.remove(SearchOptions.VECTOR)
+            item.features.remove(SearchOptions.SEMANTIC)
 
     # We need to query all nodes
     query_parser = QueryParser(
         kbid=kbid,
         features=item.features,
         query=item.query,
-        filters=item.filters,
+        label_filters=item.filters,
+        keyword_filters=[],
         faceted=item.faceted,
         sort=item.sort,
-        page_number=item.page_number,
-        page_size=item.page_size,
+        top_k=item.top_k,
         min_score=item.min_score,
         range_creation_start=item.range_creation_start,
         range_creation_end=item.range_creation_end,
@@ -452,12 +283,15 @@ async def search(
         range_modification_end=item.range_modification_end,
         fields=item.fields,
         user_vector=item.vector,
+        vectorset=item.vectorset,
         with_duplicates=item.with_duplicates,
         with_status=with_status,
         with_synonyms=item.with_synonyms,
         autofilter=item.autofilter,
         security=item.security,
         rephrase=item.rephrase,
+        hidden=await filter_hidden_resources(kbid, item.show_hidden),
+        rephrase_prompt=item.rephrase_prompt,
     )
     pb_query, incomplete_results, autofilters = await query_parser.parse()
 
@@ -470,28 +304,28 @@ async def search(
     # We need to merge
     search_results = await merge_results(
         results,
-        count=item.page_size,
-        page=item.page_number,
+        top_k=item.top_k,
         kbid=kbid,
         show=item.show,
         field_type_filter=item.field_type_filter,
         extracted=item.extracted,
-        sort=query_parser.sort,
+        sort=query_parser.sort,  # type: ignore
         requested_relations=pb_query.relation_subgraph,
         min_score=query_parser.min_score,
         highlight=item.highlight,
     )
 
     if audit is not None and do_audit:
-        await audit.search(
+        audit.search(
             kbid,
             x_nucliadb_user,
-            x_ndb_client.to_proto(),
+            to_proto.client_type(x_ndb_client),
             x_forwarded_for,
             pb_query,
             time() - start_time,
             len(search_results.resources),
         )
+
     if item.debug:
         search_results.nodes = debug_nodes_info(queried_nodes)