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

nucliadb/search/search/hydrator/__init__.py

@@ -0,0 +1,49 @@
+# 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 pydantic import BaseModel
+
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.resource import ExtractedDataTypeName
+from nucliadb_models.search import ResourceProperties
+
+
+class ResourceHydrationOptions(BaseModel):
+    """
+    Options for hydrating resources.
+    """
+
+    show: list[ResourceProperties] = []
+    extracted: list[ExtractedDataTypeName] = []
+    field_type_filter: list[FieldTypeName] = []
+
+
+class TextBlockHydrationOptions(BaseModel):
+    """
+    Options for hydrating text blocks (aka paragraphs).
+    """
+
+    # whether to highlight the text block with `<mark>...</mark>` tags or not
+    highlight: bool = False
+
+    # list of exact matches to highlight
+    ematches: list[str] | None = None
+
+    # If true, only hydrate the text block if its text is not already populated
+    only_hydrate_empty: bool = False

nucliadb/search/search/hydrator/fields.py

@@ -0,0 +1,217 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from typing import cast
+
+from typing_extensions import assert_never
+
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_NAME, FieldId
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.fields.conversation import Conversation
+from nucliadb.ingest.fields.file import File
+from nucliadb.ingest.fields.generic import Generic
+from nucliadb.ingest.fields.link import Link
+from nucliadb.ingest.fields.text import Text
+from nucliadb.models.internal.augment import ConversationProp, FieldProp, FieldText, FieldValue
+from nucliadb.search.augmentor.fields import (
+    db_augment_conversation_field,
+    db_augment_file_field,
+    db_augment_generic_field,
+    db_augment_link_field,
+    db_augment_text_field,
+)
+from nucliadb_models import hydration as hydration_models
+from nucliadb_models.common import FieldTypeName
+
+
+def page_preview_id(page_number: int) -> str:
+    """Return the string page number for an specific page"""
+    return f"{page_number}"
+
+
+async def hydrate_field(field: Field, field_id: FieldId, config: hydration_models.FieldHydration):
+    field_type = FIELD_TYPE_STR_TO_NAME[field_id.type]
+
+    if field_type == FieldTypeName.TEXT:
+        if not config.text is not None:
+            return
+        field = cast(Text, field)
+        return await hydrate_text_field(field, field_id, config.text)
+
+    elif field_type == FieldTypeName.FILE is not None:
+        if not config.file:
+            return
+        field = cast(File, field)
+        return await hydrate_file_field(field, field_id, config.file)
+
+    elif field_type == FieldTypeName.LINK is not None:
+        if not config.link:
+            return
+        field = cast(Link, field)
+        return await hydrate_link_field(field, field_id, config.link)
+
+    elif field_type == FieldTypeName.CONVERSATION is not None:
+        if not config.conversation:
+            return
+        field = cast(Conversation, field)
+        return await hydrate_conversation_field(field, field_id, config.conversation)
+
+    elif field_type == FieldTypeName.GENERIC is not None:
+        if not config.generic:
+            return
+        field = cast(Generic, field)
+        return await hydrate_generic_field(field, field_id, config.generic)
+
+    else:  # pragma: no cover
+        assert_never(field_type)
+
+
+async def hydrate_text_field(
+    field: Text,
+    field_id: FieldId,
+    config: hydration_models.TextFieldHydration,
+) -> hydration_models.HydratedTextField:
+    select: list[FieldProp] = []
+    if config.value:
+        select.append(FieldValue())
+    if config.extracted_text:
+        select.append(FieldText())
+
+    augmented = await db_augment_text_field(field, field_id, select)
+
+    hydrated = hydration_models.HydratedTextField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.TEXT,
+    )
+
+    if config.value and augmented.value:
+        hydrated.value = augmented.value
+
+    if config.extracted_text and augmented.text:
+        hydrated.extracted = hydration_models.FieldExtractedData(text=augmented.text)
+
+    return hydrated
+
+
+async def hydrate_file_field(
+    field: File,
+    field_id: FieldId,
+    config: hydration_models.FileFieldHydration,
+) -> hydration_models.HydratedFileField:
+    select: list[FieldProp] = []
+    if config.value:
+        select.append(FieldValue())
+    if config.extracted_text:
+        select.append(FieldText())
+
+    augmented = await db_augment_file_field(field, field_id, select)
+
+    hydrated = hydration_models.HydratedFileField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.FILE,
+    )
+
+    if config.value and augmented.value:
+        hydrated.value = augmented.value
+
+    if config.extracted_text and augmented.text:
+        hydrated.extracted = hydration_models.FieldExtractedData(text=augmented.text)
+
+    return hydrated
+
+
+async def hydrate_link_field(
+    field: Link,
+    field_id: FieldId,
+    config: hydration_models.LinkFieldHydration,
+) -> hydration_models.HydratedLinkField:
+    select: list[FieldProp] = []
+    if config.value:
+        select.append(FieldValue())
+    if config.extracted_text:
+        select.append(FieldText())
+
+    augmented = await db_augment_link_field(field, field_id, select)
+
+    hydrated = hydration_models.HydratedLinkField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.LINK,
+    )
+
+    if config.value and augmented.value:
+        hydrated.value = augmented.value
+
+    if config.extracted_text and augmented.text:
+        hydrated.extracted = hydration_models.FieldExtractedData(text=augmented.text)
+
+    return hydrated
+
+
+async def hydrate_conversation_field(
+    field: Conversation,
+    field_id: FieldId,
+    config: hydration_models.ConversationFieldHydration,
+) -> hydration_models.HydratedConversationField:
+    select: list[ConversationProp] = []
+    if config.value:
+        select.append(FieldValue())
+
+    augmented = await db_augment_conversation_field(field, field_id, select)
+
+    hydrated = hydration_models.HydratedConversationField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.CONVERSATION,
+    )
+
+    if config.value and augmented.value:
+        hydrated.value = augmented.value
+
+    return hydrated
+
+
+async def hydrate_generic_field(
+    field: Generic,
+    field_id: FieldId,
+    config: hydration_models.GenericFieldHydration,
+) -> hydration_models.HydratedGenericField:
+    select: list[FieldProp] = []
+    if config.value:
+        select.append(FieldValue())
+    if config.extracted_text:
+        select.append(FieldText())
+
+    augmented = await db_augment_generic_field(field, field_id, select)
+
+    hydrated = hydration_models.HydratedGenericField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.GENERIC,
+    )
+
+    if config.value and augmented.value:
+        hydrated.value = augmented.value
+
+    if config.extracted_text and augmented.text:
+        hydrated.extracted = hydration_models.FieldExtractedData(text=augmented.text)
+
+    return hydrated

nucliadb/search/search/hydrator/images.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/>.
+#
+import base64
+from typing import cast
+
+from typing_extensions import assert_never
+
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_NAME, FieldId, ParagraphId
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.fields.file import File
+from nucliadb.search import SERVICE_NAME
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.search import Image
+from nucliadb_protos import resources_pb2
+from nucliadb_utils.utilities import get_storage
+
+
+async def paragraph_source_image(
+    kbid: str, paragraph_id: ParagraphId, paragraph: resources_pb2.Paragraph
+) -> Image | None:
+    """Certain paragraphs are extracted from images using techniques like OCR or
+    inception. If that's the case, return the original image for this paragraph.
+
+    """
+    source_image = paragraph.representation.reference_file
+    if not source_image:
+        return None
+
+    if paragraph.kind not in (
+        resources_pb2.Paragraph.TypeParagraph.OCR,
+        resources_pb2.Paragraph.TypeParagraph.INCEPTION,
+    ):
+        return None
+
+    field_id = paragraph_id.field_id
+
+    # Paragraphs extracted from an image store its original image representation
+    # in the reference file. The path is incomplete though, as it's stored in
+    # the `generated` folder
+    image = await download_image(
+        kbid,
+        field_id,
+        f"generated/{source_image}",
+        # XXX: we assume all reference files are PNG images, but this actually
+        # depends on learning so it's a dangerous assumption. We should check it
+        # by ourselves
+        mime_type="image/png",
+    )
+    return image
+
+
+async def download_image(
+    kbid: str, field_id: FieldId, image_path: str, *, mime_type: str
+) -> Image | None:
+    storage = await get_storage(service_name=SERVICE_NAME)
+    sf = storage.file_extracted(
+        kbid,
+        field_id.rid,
+        field_id.type,
+        field_id.key,
+        image_path,
+    )
+    raw_image = (await storage.downloadbytes(sf.bucket, sf.key)).getvalue()
+    if not raw_image:
+        return None
+    return Image(content_type=mime_type, b64encoded=base64.b64encode(raw_image).decode())
+
+
+async def download_page_preview(field: Field, page: int) -> Image | None:
+    """Download a specific page preview for a field and return it as an Image.
+    As not all fields have previews, this function can return None.
+
+    Page previews are uploaded by learning and shared through a known path with.
+    nucliadb
+
+    """
+    field_type = FIELD_TYPE_STR_TO_NAME[field.type]
+
+    if field_type == FieldTypeName.FILE:
+        field = cast(File, field)
+        metadata = await field.get_file_extracted_data()
+
+        if metadata is None:
+            return None
+
+        assert page <= len(metadata.file_pages_previews.positions), (
+            f"paragraph page number {page} should be less or equal to the total file pages previews {len(metadata.file_pages_previews.positions)}"
+        )
+        image = await download_image(
+            field.kbid,
+            field.field_id,
+            f"generated/extracted_images_{page}.png",
+            mime_type="image/png",
+        )
+
+    elif field_type == FieldTypeName.LINK:
+        # TODO: in case of links, we want to return the link preview, that is a
+        # link converted to PDF and screenshotted
+        # REVIEW: link preview is an image or a PDF?
+        image = None
+
+    elif (
+        field_type == FieldTypeName.TEXT
+        or field_type == FieldTypeName.CONVERSATION
+        or field_type == FieldTypeName.GENERIC
+    ):
+        # these fields don't have previews
+        image = None
+
+    else:  # pragma: no cover
+        assert_never(field_type)
+
+    return image