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

nucliadb/search/augmentor/metrics.py

@@ -0,0 +1,24 @@
+# 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 nucliadb_telemetry.metrics import Observer
+
+augmentor_observer = Observer("augmentor", labels={"type": ""})

nucliadb/search/augmentor/paragraphs.py

@@ -0,0 +1,334 @@
+# 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 asyncio
+from collections.abc import Sequence
+from typing import cast
+
+from typing_extensions import assert_never
+
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB, ParagraphId
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.models.internal.augment import (
+    AugmentedParagraph,
+    AugmentedRelatedParagraphs,
+    Metadata,
+    Paragraph,
+    ParagraphImage,
+    ParagraphPage,
+    ParagraphPosition,
+    ParagraphProp,
+    ParagraphTable,
+    ParagraphText,
+    RelatedParagraphs,
+)
+from nucliadb.search.augmentor.metrics import augmentor_observer
+from nucliadb.search.augmentor.utils import limited_concurrency
+from nucliadb.search.search import cache
+from nucliadb.search.search.paragraphs import get_paragraph_from_full_text
+from nucliadb_models.search import TextPosition
+from nucliadb_protos import resources_pb2
+
+
+async def augment_paragraphs(
+    kbid: str,
+    given: list[Paragraph],
+    select: list[ParagraphProp],
+    *,
+    concurrency_control: asyncio.Semaphore | None = None,
+) -> dict[ParagraphId, AugmentedParagraph | None]:
+    """Augment a list of paragraphs following an augmentation"""
+
+    ops = []
+    for paragraph in given:
+        task = asyncio.create_task(
+            limited_concurrency(
+                augment_paragraph(kbid, paragraph.id, select, paragraph.metadata),
+                max_ops=concurrency_control,
+            )
+        )
+        ops.append(task)
+    results: list[AugmentedParagraph | None] = await asyncio.gather(*ops)
+
+    augmented = {}
+    for paragraph, augmentation in zip(given, results):
+        augmented[paragraph.id] = augmentation
+
+    return augmented
+
+
+async def augment_paragraph(
+    kbid: str,
+    paragraph_id: ParagraphId,
+    select: list[ParagraphProp],
+    metadata: Metadata | None,
+) -> AugmentedParagraph | None:
+    rid = paragraph_id.rid
+    resource = await cache.get_resource(kbid, rid)
+    if resource is None:
+        # skip resources that aren't in the DB
+        return None
+
+    field_id = paragraph_id.field_id
+    field_type_pb = FIELD_TYPE_STR_TO_PB[field_id.type]
+    # we must check if field exists or get_field will return an empty field
+    # (behaviour thought for ingestion) that we don't want
+    if not (await resource.field_exists(field_type_pb, field_id.key)):
+        # skip a fields that aren't in the DB
+        return None
+    field = await resource.get_field(field_id.key, field_id.pb_type)
+
+    return await db_augment_paragraph(resource, field, paragraph_id, select, metadata)
+
+
+async def db_augment_paragraph(
+    resource: Resource,
+    field: Field,
+    paragraph_id: ParagraphId,
+    select: list[ParagraphProp],
+    metadata: Metadata | None,
+) -> AugmentedParagraph:
+    select = dedup_paragraph_select(select)
+
+    # we use an accessor to get the metadata to avoid unnecessary DB round
+    # trips. With this, we'll only fetch it one and only if we need it
+    _metadata = metadata
+    _metadata_available = True
+
+    async def access_metadata() -> Metadata | None:
+        nonlocal _metadata, _metadata_available
+
+        if _metadata is None and _metadata_available:
+            _metadata = await db_paragraph_metadata(field, paragraph_id)
+
+        if _metadata is None:
+            _metadata_available = False
+
+        return _metadata
+
+    text = None
+    position = None
+    image_path = None
+    table_path = None
+    page_preview_path = None
+    related = None
+    for prop in select:
+        if isinstance(prop, ParagraphText):
+            text = await get_paragraph_text(field, paragraph_id)
+
+        elif isinstance(prop, ParagraphPosition):
+            position = await get_paragraph_position(field, paragraph_id)
+
+        elif isinstance(prop, ParagraphImage):
+            metadata = await access_metadata()
+            if metadata is None:
+                continue
+            if metadata.is_an_image and metadata.source_file:
+                image_path = f"generated/{metadata.source_file}"
+
+        elif isinstance(prop, ParagraphTable):
+            metadata = await access_metadata()
+            if metadata is None:
+                continue
+            if metadata.is_a_table:
+                if prop.prefer_page_preview and metadata.page and metadata.in_page_with_visual:
+                    page_preview_path = f"generated/extracted_images_{metadata.page}.png"
+                    table_path = page_preview_path
+                elif metadata.source_file:
+                    image_path = f"generated/{metadata.source_file}"
+                    table_path = image_path
+
+        elif isinstance(prop, ParagraphPage):
+            if prop.preview:
+                metadata = await access_metadata()
+                if metadata is None:
+                    continue
+                if metadata.page and metadata.in_page_with_visual:
+                    page_preview_path = f"generated/extracted_images_{metadata.page}.png"
+
+        elif isinstance(prop, RelatedParagraphs):
+            related = await related_paragraphs(
+                field,
+                paragraph_id,
+                neighbours_before=prop.neighbours_before,
+                neighbours_after=prop.neighbours_after,
+            )
+
+        else:  # pragma: no cover
+            assert_never(prop)
+
+    return AugmentedParagraph(
+        id=paragraph_id,
+        text=text,
+        position=position,
+        source_image_path=image_path,
+        table_image_path=table_path,
+        page_preview_path=page_preview_path,
+        related=related,
+    )
+
+
+def dedup_paragraph_select(select: list[ParagraphProp]) -> list[ParagraphProp]:
+    """Merge any duplicated property taking the broader augmentation possible."""
+    merged = {}
+    for prop in select:
+        if prop.prop not in merged:
+            merged[prop.prop] = prop
+
+        else:
+            m = merged[prop.prop]
+
+            if (
+                isinstance(prop, ParagraphText)
+                or isinstance(prop, ParagraphPosition)
+                or isinstance(prop, ParagraphImage)
+            ):
+                # properties without parameters
+                pass
+
+            elif isinstance(prop, ParagraphTable):
+                prop = cast(ParagraphTable, prop)
+                m = cast(ParagraphTable, m)
+                m.prefer_page_preview = m.prefer_page_preview or prop.prefer_page_preview
+
+            elif isinstance(prop, ParagraphPage):
+                prop = cast(ParagraphPage, prop)
+                m = cast(ParagraphPage, m)
+                m.preview = m.preview or prop.preview
+
+            elif isinstance(prop, RelatedParagraphs):
+                prop = cast(RelatedParagraphs, prop)
+                m = cast(RelatedParagraphs, m)
+                m.neighbours_before = max(m.neighbours_before, prop.neighbours_before)
+                m.neighbours_after = max(m.neighbours_after, prop.neighbours_after)
+
+            else:  # pragma: no cover
+                assert_never(prop)
+
+    return list(merged.values())
+
+
+async def db_paragraph_metadata(field: Field, paragraph_id: ParagraphId) -> Metadata | None:
+    """Obtain paragraph metadata from the source of truth (maindb/blob).
+
+    This operation may require data from blob storage, which makes it costly.
+
+    """
+    field_paragraphs = await get_field_paragraphs(field)
+    if field_paragraphs is None:
+        # We don't have paragraph metadata for this field, we can't do anything
+        return None
+
+    for paragraph in field_paragraphs:
+        field_paragraph_id = field.field_id.paragraph_id(paragraph.start, paragraph.end)
+        if field_paragraph_id == paragraph_id:
+            metadata = Metadata.from_db_paragraph(paragraph)
+            return metadata
+    else:
+        return None
+
+
+async def get_field_paragraphs(field: Field) -> Sequence[resources_pb2.Paragraph] | None:
+    field_metadata = await field.get_field_metadata()
+    if field_metadata is None:
+        return None
+
+    field_id = field.field_id
+    if field_id.subfield_id is None:
+        field_paragraphs = field_metadata.metadata.paragraphs
+    else:
+        field_paragraphs = field_metadata.split_metadata[field_id.subfield_id].paragraphs
+
+    return field_paragraphs
+
+
+@augmentor_observer.wrap({"type": "paragraph_text"})
+async def get_paragraph_text(field: Field, paragraph_id: ParagraphId) -> str | None:
+    text = await get_paragraph_from_full_text(
+        field=field,
+        start=paragraph_id.paragraph_start,
+        end=paragraph_id.paragraph_end,
+        split=paragraph_id.field_id.subfield_id,
+        log_on_missing_field=True,
+    )
+    # we want to be explicit with not having the paragraph text but the function
+    # above returns an empty string if it can't find it
+    return text or None
+
+
+async def get_paragraph_position(field: Field, paragraph_id: ParagraphId) -> TextPosition | None:
+    field_paragraphs = await get_field_paragraphs(field)
+    if field_paragraphs is None:
+        return None
+
+    idx: int | None
+    for idx, paragraph in enumerate(field_paragraphs):
+        field_paragraph_id = field.field_id.paragraph_id(paragraph.start, paragraph.end)
+        if field_paragraph_id == paragraph_id:
+            break
+    else:
+        # we haven't found the paragraph, we can't provide a position
+        return None
+
+    return TextPosition(
+        index=idx,
+        start=paragraph.start,
+        end=paragraph.end,
+        start_seconds=list(paragraph.start_seconds),
+        end_seconds=list(paragraph.end_seconds),
+    )
+
+
+async def related_paragraphs(
+    field: Field,
+    paragraph_id: ParagraphId,
+    *,
+    neighbours_before: int = 0,
+    neighbours_after: int = 0,
+) -> AugmentedRelatedParagraphs | None:
+    field_paragraphs = await get_field_paragraphs(field)
+    if field_paragraphs is None:
+        return None
+
+    idx: int | None
+    for idx, paragraph in enumerate(field_paragraphs):
+        field_paragraph_id = field.field_id.paragraph_id(paragraph.start, paragraph.end)
+        if field_paragraph_id == paragraph_id:
+            break
+    else:
+        # we haven't found the paragraph, we won't find any related either
+        return None
+
+    before = []
+    for idx_before in range(max(idx - neighbours_before, 0), idx):
+        paragraph = field_paragraphs[idx_before]
+        paragraph_id = field.field_id.paragraph_id(paragraph.start, paragraph.end)
+        before.append(paragraph_id)
+
+    after = []
+    for idx_after in range(idx + 1, min(idx + 1 + neighbours_after, len(field_paragraphs))):
+        paragraph = field_paragraphs[idx_after]
+        paragraph_id = field.field_id.paragraph_id(paragraph.start, paragraph.end)
+        after.append(paragraph_id)
+
+    return AugmentedRelatedParagraphs(
+        neighbours_before=before,
+        neighbours_after=after,
+    )

nucliadb/search/augmentor/resources.py

@@ -0,0 +1,238 @@
+# 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 asyncio
+
+from typing_extensions import assert_never
+
+import nucliadb_models.resource
+from nucliadb.common import datamanagers
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.ingest.serialize import (
+    serialize_extra,
+    serialize_origin,
+    serialize_resource,
+    serialize_security,
+)
+from nucliadb.models.internal.augment import (
+    AugmentedResource,
+    ResourceClassificationLabels,
+    ResourceExtra,
+    ResourceOrigin,
+    ResourceProp,
+    ResourceSecurity,
+    ResourceSummary,
+    ResourceTitle,
+)
+from nucliadb.search.augmentor.metrics import augmentor_observer
+from nucliadb.search.augmentor.utils import limited_concurrency
+from nucliadb.search.search import cache
+from nucliadb.search.search.hydrator import ResourceHydrationOptions
+from nucliadb_models.search import ResourceProperties
+from nucliadb_protos import resources_pb2
+from nucliadb_utils import const
+from nucliadb_utils.utilities import has_feature
+
+
+async def augment_resources(
+    kbid: str,
+    given: list[str],
+    select: list[ResourceProp],
+    *,
+    concurrency_control: asyncio.Semaphore | None = None,
+) -> dict[str, AugmentedResource | None]:
+    """Augment a list of resources following an augmentation"""
+
+    ops = []
+    for rid in given:
+        task = asyncio.create_task(
+            limited_concurrency(
+                augment_resource(kbid, rid, select),
+                max_ops=concurrency_control,
+            )
+        )
+        ops.append(task)
+    results: list[AugmentedResource | None] = await asyncio.gather(*ops)
+
+    augmented = {}
+    for rid, augmentation in zip(given, results):
+        augmented[rid] = augmentation
+
+    return augmented
+
+
+async def augment_resource(
+    kbid: str,
+    rid: str,
+    select: list[ResourceProp],
+) -> AugmentedResource | None:
+    resource = await cache.get_resource(kbid, rid)
+    if resource is None:
+        # skip resources that aren't in the DB
+        return None
+
+    return await db_augment_resource(resource, select)
+
+
+@augmentor_observer.wrap({"type": "db_resource"})
+async def db_augment_resource(
+    resource: Resource,
+    select: list[ResourceProp],
+) -> AugmentedResource:
+    select = dedup_resource_select(select)
+
+    title = None
+    summary = None
+    origin = None
+    extra = None
+    security = None
+    labels = None
+
+    basic = None
+    for prop in select:
+        if isinstance(prop, ResourceTitle):
+            if basic is None:
+                basic = await resource.get_basic()
+            if basic is not None:
+                title = basic.title
+
+        elif isinstance(prop, ResourceSummary):
+            if basic is None:
+                basic = await resource.get_basic()
+            if basic is not None:
+                summary = basic.summary
+
+        elif isinstance(prop, ResourceOrigin):
+            origin = await serialize_origin(resource)
+
+        elif isinstance(prop, ResourceExtra):
+            extra = await serialize_extra(resource)
+
+        elif isinstance(prop, ResourceSecurity):
+            security = await serialize_security(resource)
+
+        elif isinstance(prop, ResourceClassificationLabels):
+            labels = await classification_labels(resource)
+
+        else:
+            assert_never(prop)
+
+    augmented = AugmentedResource(
+        id=resource.uuid,
+        title=title,
+        summary=summary,
+        origin=origin,
+        extra=extra,
+        security=security,
+        classification_labels=labels,
+    )
+    return augmented
+
+
+def dedup_resource_select(select: list[ResourceProp]) -> list[ResourceProp]:
+    # there's no resource prop with fields that need special treatement to
+    # merge, just get by unique prop id
+    merged: dict[str, ResourceProp] = {}
+    for prop in select:
+        merged.setdefault(prop.prop, prop)
+    return list(merged.values())
+
+
+async def get_basic(resource: Resource) -> resources_pb2.Basic | None:
+    # HACK: resource.get_basic() always returns a pb, even if it's not in the
+    # DB. Here we really want to know if there's basic or not
+    basic = await datamanagers.resources.get_basic(resource.txn, kbid=resource.kbid, rid=resource.uuid)
+    return basic
+
+
+async def classification_labels(resource: Resource) -> dict[str, set[str]] | None:
+    basic = await get_basic(resource)
+    if basic is None:
+        return None
+
+    labels: dict[str, set[str]] = {}
+    for classification in basic.usermetadata.classifications:
+        labels.setdefault(classification.labelset, set()).add(classification.label)
+    return labels
+
+
+async def augment_resources_deep(
+    kbid: str,
+    given: list[str],
+    opts: ResourceHydrationOptions,
+    *,
+    concurrency_control: asyncio.Semaphore | None = None,
+) -> dict[str, nucliadb_models.resource.Resource | None]:
+    """Augment resources using the Resource model. Depending on the options,
+    this can serialize resource fields, extracted data like text, vectors...
+
+    Thus, this operation can be quite expensive.
+
+    """
+
+    if ResourceProperties.EXTRACTED in opts.show and has_feature(
+        const.Features.IGNORE_EXTRACTED_IN_SEARCH, context={"kbid": kbid}, default=False
+    ):
+        # Returning extracted metadata in search results is deprecated and this flag
+        # will be set to True for all KBs in the future.
+        opts.show.remove(ResourceProperties.EXTRACTED)
+        opts.extracted.clear()
+
+    ops = []
+    for rid in given:
+        task = asyncio.create_task(
+            limited_concurrency(
+                augment_resource_deep(kbid, rid, opts),
+                max_ops=concurrency_control,
+            )
+        )
+        ops.append(task)
+    results: list[nucliadb_models.resource.Resource | None] = await asyncio.gather(*ops)
+
+    augmented: dict[str, nucliadb_models.resource.Resource | None] = {}
+    for rid, augmentation in zip(given, results):
+        augmented[rid] = augmentation
+
+    return augmented
+
+
+@augmentor_observer.wrap({"type": "seialize_resource"})
+async def augment_resource_deep(
+    kbid: str,
+    rid: str,
+    opts: ResourceHydrationOptions,
+) -> nucliadb_models.resource.Resource | None:
+    """Augment a resource using the Resource model. Depending on the options,
+    this can serialize resource fields, extracted data like text, vectors...
+
+    Thus, this operation can be quite expensive.
+
+    """
+    resource = await cache.get_resource(kbid, rid)
+    if resource is None:
+        # skip resources that aren't in the DB
+        return None
+
+    serialized = await serialize_resource(
+        resource,
+        show=opts.show,
+        field_type_filter=opts.field_type_filter,
+        extracted=opts.extracted,
+    )
+    return serialized

nucliadb/search/augmentor/utils.py

@@ -0,0 +1,33 @@
+# 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 asyncio
+from collections.abc import Awaitable
+from contextlib import AsyncExitStack
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+async def limited_concurrency(aw: Awaitable[T], *, max_ops: asyncio.Semaphore | None) -> T:
+    async with AsyncExitStack() as stack:
+        if max_ops is not None:
+            await stack.enter_async_context(max_ops)
+        r = await aw
+    return r

nucliadb/search/lifecycle.py

@@ -24,7 +24,7 @@ from fastapi import FastAPI
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.context.fastapi import inject_app_context
 from nucliadb.common.maindb.utils import setup_driver
-from nucliadb.common.nidx import start_nidx_utility
+from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb.ingest.utils import start_ingest, stop_ingest
 from nucliadb.search import SERVICE_NAME
 from nucliadb.search.predict import start_predict_engine
@@ -61,6 +61,8 @@ async def lifespan(app: FastAPI):
     if get_utility(Utility.PREDICT):
         clean_utility(Utility.PREDICT)
 
+    await stop_nidx_utility()
+
     await finalize_utilities()
     await stop_audit_utility()
     await teardown_cluster()