nucliadb 6.8.1.post4983__py3-none-any.whl → 6.8.1.post5003__py3-none-any.whl

nucliadb/common/ids.py

@@ -47,6 +47,8 @@ FIELD_TYPE_NAME_TO_STR = {
     FieldTypeName.CONVERSATION: "c",
 }
 
+FIELD_TYPE_STR_TO_NAME = {v: k for k, v in FIELD_TYPE_NAME_TO_STR.items()}
+
 
 @dataclass
 class FieldId:
@@ -65,7 +67,7 @@ class FieldId:
 
     Examples:
 
-    >>> FieldId(rid="rid", type="u", key="/my-link")
+    >>> FieldId(rid="rid", type="u", key="my-link")
     FieldID("rid/u/my-link")
     >>> FieldId.from_string("rid/u/my-link")
     FieldID("rid/u/my-link")
@@ -77,31 +79,6 @@ class FieldId:
     # also knwon as `split`, this indicates a part of a field in, for example, conversations
     subfield_id: Optional[str] = None
 
-    def __repr__(self) -> str:
-        return f"FieldId({self.full()})"
-
-    def short_without_subfield(self) -> str:
-        return f"/{self.type}/{self.key}"
-
-    def full(self) -> str:
-        if self.subfield_id is None:
-            return f"{self.rid}/{self.type}/{self.key}"
-        else:
-            return f"{self.rid}/{self.type}/{self.key}/{self.subfield_id}"
-
-    def __hash__(self) -> int:
-        return hash(self.full())
-
-    @property
-    def pb_type(self) -> FieldType.ValueType:
-        return FIELD_TYPE_STR_TO_PB[self.type]
-
-    @classmethod
-    def from_pb(
-        cls, rid: str, field_type: FieldType.ValueType, key: str, subfield_id: Optional[str] = None
-    ) -> "FieldId":
-        return cls(rid=rid, type=FIELD_TYPE_PB_TO_STR[field_type], key=key, subfield_id=subfield_id)
-
     @classmethod
     def from_string(cls, value: str) -> "FieldId":
         """
@@ -120,11 +97,11 @@ class FieldId:
         parts = value.split("/")
         if len(parts) == 3:
             rid, _type, key = parts
-            _type = cls.parse_field_type(_type)
+            _type = cls._parse_field_type(_type)
             return cls(rid=rid, type=_type, key=key)
         elif len(parts) == 4:
             rid, _type, key, subfield_id = parts
-            _type = cls.parse_field_type(_type)
+            _type = cls._parse_field_type(_type)
             return cls(
                 rid=rid,
                 type=_type,
@@ -135,7 +112,46 @@ class FieldId:
             raise ValueError(f"Invalid FieldId: {value}")
 
     @classmethod
-    def parse_field_type(cls, _type: str) -> str:
+    def from_pb(
+        cls, rid: str, field_type: FieldType.ValueType, key: str, subfield_id: Optional[str] = None
+    ) -> "FieldId":
+        return cls(rid=rid, type=FIELD_TYPE_PB_TO_STR[field_type], key=key, subfield_id=subfield_id)
+
+    @property
+    def pb_type(self) -> FieldType.ValueType:
+        return FIELD_TYPE_STR_TO_PB[self.type]
+
+    def full(self) -> str:
+        if self.subfield_id is None:
+            return f"{self.rid}/{self.type}/{self.key}"
+        else:
+            return f"{self.rid}/{self.type}/{self.key}/{self.subfield_id}"
+
+    def short_without_subfield(self) -> str:
+        return f"/{self.type}/{self.key}"
+
+    def paragraph_id(self, paragraph_start: int, paragraph_end: int) -> "ParagraphId":
+        """Generate a ParagraphId from the current field given its start and
+        end.
+
+        """
+        return ParagraphId(
+            field_id=self,
+            paragraph_start=paragraph_start,
+            paragraph_end=paragraph_end,
+        )
+
+    def __str__(self) -> str:
+        return self.full()
+
+    def __repr__(self) -> str:
+        return f"FieldId({self.full()})"
+
+    def __hash__(self) -> int:
+        return hash(self.full())
+
+    @staticmethod
+    def _parse_field_type(_type: str) -> str:
         if _type not in FIELD_TYPE_STR_TO_PB:
             # Try to parse the enum value
             # XXX: This is to support field types that are integer values of FieldType
@@ -157,19 +173,6 @@ class ParagraphId:
     paragraph_start: int
     paragraph_end: int
 
-    def __repr__(self) -> str:
-        return f"ParagraphId({self.full()})"
-
-    def full(self) -> str:
-        return f"{self.field_id.full()}/{self.paragraph_start}-{self.paragraph_end}"
-
-    def __hash__(self) -> int:
-        return hash(self.full())
-
-    @property
-    def rid(self) -> str:
-        return self.field_id.rid
-
     @classmethod
     def from_string(cls, value: str) -> "ParagraphId":
         parts = value.split("/")
@@ -192,6 +195,22 @@ class ParagraphId:
             paragraph_end=vid.vector_end,
         )
 
+    @property
+    def rid(self) -> str:
+        return self.field_id.rid
+
+    def full(self) -> str:
+        return f"{self.field_id.full()}/{self.paragraph_start}-{self.paragraph_end}"
+
+    def __str__(self) -> str:
+        return self.full()
+
+    def __repr__(self) -> str:
+        return f"ParagraphId({self.full()})"
+
+    def __hash__(self) -> int:
+        return hash(self.full())
+
 
 @dataclass
 class VectorId:
@@ -217,19 +236,6 @@ class VectorId:
     vector_start: int
     vector_end: int
 
-    def __repr__(self) -> str:
-        return f"VectorId({self.full()})"
-
-    def full(self) -> str:
-        return f"{self.field_id.full()}/{self.index}/{self.vector_start}-{self.vector_end}"
-
-    def __hash__(self) -> int:
-        return hash(self.full())
-
-    @property
-    def rid(self) -> str:
-        return self.field_id.rid
-
     @classmethod
     def from_string(cls, value: str) -> "VectorId":
         parts = value.split("/")
@@ -239,6 +245,22 @@ class VectorId:
         field_id = FieldId.from_string("/".join(parts[:-2]))
         return cls(field_id=field_id, index=index, vector_start=start, vector_end=end)
 
+    @property
+    def rid(self) -> str:
+        return self.field_id.rid
+
+    def full(self) -> str:
+        return f"{self.field_id.full()}/{self.index}/{self.vector_start}-{self.vector_end}"
+
+    def __str__(self) -> str:
+        return self.full()
+
+    def __repr__(self) -> str:
+        return f"VectorId({self.full()})"
+
+    def __hash__(self) -> int:
+        return hash(self.full())
+
 
 def extract_data_augmentation_id(generated_field_id: str) -> Optional[str]:
     """Data augmentation generated fields have a strict id with the following

nucliadb/ingest/fields/base.py

@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING, Any, Generic, Optional, Type, TypeVar
 from google.protobuf.message import DecodeError, Message
 
 from nucliadb.common import datamanagers
+from nucliadb.common.ids import FieldId
 from nucliadb.ingest.fields.exceptions import InvalidFieldClass, InvalidPBClass
 from nucliadb_protos.knowledgebox_pb2 import VectorSetConfig
 from nucliadb_protos.resources_pb2 import (
@@ -125,6 +126,14 @@ class Field(Generic[PbType]):
     def uuid(self) -> str:
         return self.resource.uuid
 
+    @property
+    def field_id(self) -> FieldId:
+        return FieldId(
+            rid=self.resource.uuid,
+            type=self.type,
+            key=self.id,
+        )
+
     @property
     def storage(self) -> Storage:
         return self.resource.storage

nucliadb/ingest/orm/resource.py

@@ -135,7 +135,7 @@ class Resource:
         await self.txn.set(new_key, self.uuid.encode())
 
     # Basic
-    async def get_basic(self) -> Optional[PBBasic]:
+    async def get_basic(self) -> PBBasic:
         if self.basic is None:
             basic = await datamanagers.resources.get_basic(self.txn, kbid=self.kb.kbid, rid=self.uuid)
             self.basic = basic if basic is not None else PBBasic()
@@ -354,7 +354,16 @@ class Resource:
 
         await field_obj.delete()
 
+    async def field_exists(self, type: FieldType.ValueType, field: str) -> bool:
+        """Return whether this resource has this field or not."""
+        all_fields_ids = await self.get_fields_ids()
+        for field_type, field_id in all_fields_ids:
+            if field_type == type and field_id == field:
+                return True
+        return False
+
     def has_field(self, type: FieldType.ValueType, field: str) -> bool:
+        # REVIEW: are we sure we don't want to actually check this?
         return (type, field) in self.fields
 
     async def get_all_field_ids(self, *, for_update: bool) -> Optional[PBAllFieldIDs]:

nucliadb/search/api/v1/__init__.py

@@ -23,6 +23,7 @@ from . import (  # noqa: F401
     feedback,
     find,
     graph,
+    hydrate,
     knowledgebox,
     predict_proxy,
     search,

nucliadb/search/api/v1/hydrate.py

@@ -0,0 +1,328 @@
+# 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 import Awaitable, Optional, Union
+
+from async_lru import alru_cache
+from fastapi import Request, Response
+from fastapi_versioning import version
+
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB, FieldId, ParagraphId
+from nucliadb.ingest.fields.base import Field
+from nucliadb.search.api.v1.router import KB_PREFIX, api
+from nucliadb.search.search import cache
+from nucliadb.search.search.cache import request_caches
+from nucliadb.search.search.hydrator.fields import hydrate_field, page_preview_id
+from nucliadb.search.search.hydrator.images import (
+    download_page_preview,
+)
+from nucliadb.search.search.hydrator.paragraphs import ParagraphIndex, hydrate_paragraph
+from nucliadb.search.search.hydrator.resources import hydrate_resource
+from nucliadb_models.hydration import (
+    Hydrated,
+    HydratedConversationField,
+    HydratedFileField,
+    HydratedGenericField,
+    HydratedLinkField,
+    HydratedParagraph,
+    HydratedResource,
+    HydratedTextField,
+    HydrateRequest,
+    Hydration,
+    ParagraphHydration,
+)
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.search import Image
+from nucliadb_utils.authentication import requires
+
+
+@api.post(
+    f"/{KB_PREFIX}/{{kbid}}/hydrate",
+    status_code=200,
+    summary="Hydrate a set of paragraphs",
+    description="Internal API endpoint to hydrate a set of paragraphs",
+    include_in_schema=False,
+    response_model_exclude_unset=True,
+    tags=["Hydration"],
+)
+@requires(NucliaDBRoles.READER)
+@version(1)
+async def hydrate_endpoint(
+    request: Request,
+    response: Response,
+    kbid: str,
+    item: HydrateRequest,
+) -> Hydrated:
+    with request_caches():
+        return await Hydrator(kbid, item.hydration).hydrate(item.data)
+
+
+class HydratedBuilder:
+    """Builder class to construct an Hydrated payload."""
+
+    def __init__(self) -> None:
+        self._resources: dict[str, HydratedResource] = {}
+        self._fields: dict[
+            str,
+            Union[
+                HydratedTextField,
+                HydratedFileField,
+                HydratedLinkField,
+                HydratedConversationField,
+                HydratedGenericField,
+            ],
+        ] = {}
+        self._paragraphs: dict[str, HydratedParagraph] = {}
+
+    @property
+    def resources(self) -> dict[str, HydratedResource]:
+        return self._resources
+
+    @property
+    def fields(
+        self,
+    ) -> dict[
+        str,
+        Union[
+            HydratedTextField,
+            HydratedFileField,
+            HydratedLinkField,
+            HydratedConversationField,
+            HydratedGenericField,
+        ],
+    ]:
+        return self._fields
+
+    @property
+    def paragraphs(self) -> dict[str, HydratedParagraph]:
+        return self._paragraphs
+
+    def build(self) -> Hydrated:
+        return Hydrated(
+            resources=self._resources,
+            fields=self._fields,
+            paragraphs=self._paragraphs,
+        )
+
+    def add_resource(self, rid: str, resource: HydratedResource):
+        self._resources[rid] = resource
+
+    def add_field(
+        self,
+        field_id: FieldId,
+        field: Union[
+            HydratedTextField,
+            HydratedFileField,
+            HydratedLinkField,
+            HydratedConversationField,
+            HydratedGenericField,
+        ],
+    ):
+        self._fields[field_id.full()] = field
+
+    def has_field(self, field_id: FieldId) -> bool:
+        return field_id.full() in self._fields
+
+    def add_paragraph(self, paragraph_id: ParagraphId, paragraph: HydratedParagraph):
+        self._paragraphs[paragraph_id.full()] = paragraph
+
+    def add_page_preview(self, paragraph_id: ParagraphId, page: int, image: Image):
+        field_id = paragraph_id.field_id
+        field = self._fields[field_id.full()]
+
+        if not isinstance(field, HydratedFileField):
+            # Other field types have no page preview concept
+            return
+
+        if field.previews is None:
+            field.previews = {}
+
+        preview_id = page_preview_id(page)
+        field.previews[preview_id] = image
+
+        paragraph = self._paragraphs[paragraph_id.full()]
+        assert paragraph.page is not None, "should already be set"
+        paragraph.page.page_preview_ref = preview_id
+
+    def add_table_page_preview(self, paragraph_id: ParagraphId, page: int, image: Image):
+        field_id = paragraph_id.field_id
+        field = self._fields[field_id.full()]
+
+        if not isinstance(field, HydratedFileField):
+            # Other field types have no page preview concept
+            return
+
+        if field.previews is None:
+            field.previews = {}
+
+        preview_id = page_preview_id(page)
+        field.previews[preview_id] = image
+
+        paragraph = self._paragraphs[paragraph_id.full()]
+        assert paragraph.table is not None, "should already be set"
+        paragraph.table.page_preview_ref = preview_id
+
+
+class Hydrator:
+    def __init__(self, kbid: str, config: Hydration):
+        self.kbid = kbid
+        self.config = config
+        self.hydrated = HydratedBuilder()
+
+        # cached paragraphs per field
+        self.field_paragraphs: dict[FieldId, ParagraphIndex] = {}
+
+        self.max_ops = asyncio.Semaphore(50)
+
+    async def hydrate(self, paragraph_ids: list[str]) -> Hydrated:
+        paragraph_tasks = {}
+        field_tasks = {}
+        resource_tasks = {}
+
+        unique_paragraph_ids = set(paragraph_ids)
+        for user_paragraph_id in unique_paragraph_ids:
+            try:
+                paragraph_id = ParagraphId.from_string(user_paragraph_id)
+            except ValueError:
+                # skip paragraphs with invalid format
+                continue
+
+            field_id = paragraph_id.field_id
+            rid = paragraph_id.rid
+
+            resource = await cache.get_resource(self.kbid, rid)
+            if resource is None:
+                # skip resources that aren't in the DB
+                continue
+
+            field_type_pb = FIELD_TYPE_STR_TO_PB[field_id.type]
+            if not (await resource.field_exists(field_type_pb, field_id.key)):
+                # skip a fields that aren't in the DB
+                continue
+            field = await resource.get_field(field_id.key, field_id.pb_type)
+
+            if field_id not in self.field_paragraphs:
+                field_paragraphs_index = ParagraphIndex(field_id)
+                self.field_paragraphs[field_id] = field_paragraphs_index
+            field_paragraphs_index = self.field_paragraphs[field_id]
+
+            paragraph_tasks[paragraph_id] = asyncio.create_task(
+                self._limited_concurrency(
+                    hydrate_paragraph(
+                        resource, field, paragraph_id, self.config.paragraph, field_paragraphs_index
+                    ),
+                )
+            )
+
+            if field_id not in field_tasks:
+                field_tasks[field_id] = asyncio.create_task(
+                    self._limited_concurrency(hydrate_field(resource, field_id, self.config.field))
+                )
+
+            if rid not in resource_tasks:
+                if self.config.resource is not None:
+                    resource_tasks[rid] = asyncio.create_task(
+                        self._limited_concurrency(hydrate_resource(resource, rid, self.config.resource))
+                    )
+
+        ops = [
+            *paragraph_tasks.values(),
+            *field_tasks.values(),
+            *resource_tasks.values(),
+        ]
+        results = await asyncio.gather(*ops)
+        hydrated_paragraphs = results[: len(paragraph_tasks)]
+        hydrated_fields = results[len(paragraph_tasks) : len(paragraph_tasks) + len(field_tasks)]
+        hydrated_resources = results[
+            len(paragraph_tasks) + len(field_tasks) : len(paragraph_tasks)
+            + len(field_tasks)
+            + len(resource_tasks)
+        ]
+
+        for rid, hydrated_resource in zip(resource_tasks.keys(), hydrated_resources):
+            self.hydrated.add_resource(rid, hydrated_resource)
+
+        for field_id, hydrated_field in zip(field_tasks.keys(), hydrated_fields):
+            if hydrated_field is not None:
+                self.hydrated.add_field(field_id, hydrated_field)
+
+        for paragraph_id, (hydrated_paragraph, extra) in zip(
+            paragraph_tasks.keys(), hydrated_paragraphs
+        ):
+            self.hydrated.add_paragraph(paragraph_id, hydrated_paragraph)
+
+            for related_paragraph_id in extra.related_paragraph_ids:
+                field_id = related_paragraph_id.field_id
+                rid = related_paragraph_id.rid
+
+                resource = await cache.get_resource(self.kbid, rid)
+                if resource is None:
+                    # skip resources that aren't in the DB
+                    continue
+
+                field_type_pb = FIELD_TYPE_STR_TO_PB[field_id.type]
+                if not (await resource.field_exists(field_type_pb, field_id.key)):
+                    # skip a fields that aren't in the DB
+                    continue
+                field = await resource.get_field(field_id.key, field_id.pb_type)
+
+                if field_id not in self.field_paragraphs:
+                    field_paragraphs_index = ParagraphIndex(field_id)
+                    self.field_paragraphs[field_id] = field_paragraphs_index
+                field_paragraphs_index = self.field_paragraphs[field_id]
+
+                (hydrated_paragraph, _) = await hydrate_paragraph(
+                    resource,
+                    field,
+                    related_paragraph_id,
+                    ParagraphHydration(
+                        text=self.config.paragraph.text, image=None, table=None, page=None, related=None
+                    ),
+                    field_paragraphs_index,
+                )
+                self.hydrated.add_paragraph(related_paragraph_id, hydrated_paragraph)
+
+            if self.hydrated.has_field(field_id):
+                # we only hydrate page and table previews for fields the user
+                # allowed hydration, skipping fields with explicitly disabled
+                # hydration
+
+                if extra.field_page is not None:
+                    page_number = extra.field_page
+                    preview = await self.cached_download_page_preview(field, page_number)
+                    if preview is not None:
+                        self.hydrated.add_page_preview(paragraph_id, page_number, preview)
+
+                if extra.field_table_page is not None:
+                    page_number = extra.field_table_page
+                    preview = await self.cached_download_page_preview(field, page_number)
+                    if preview is not None:
+                        self.hydrated.add_table_page_preview(paragraph_id, page_number, preview)
+
+        return self.hydrated.build()
+
+    # TODO: proper typing
+    async def _limited_concurrency(self, aw: Awaitable):
+        async with self.max_ops:
+            return await aw
+
+    @alru_cache(maxsize=None)
+    async def cached_download_page_preview(self, field: Field, page: int) -> Optional[Image]:
+        return await download_page_preview(field, page)

nucliadb/search/search/cache.py

@@ -21,8 +21,6 @@ import contextlib
 import logging
 from typing import Optional
 
-import backoff
-
 from nucliadb.common.cache import (
     extracted_text_cache,
     get_extracted_text_cache,
@@ -74,23 +72,6 @@ async def get_field_extracted_text(field: Field) -> Optional[ExtractedText]:
     return extracted_text
 
 
-@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
-async def field_get_extracted_text(field: Field) -> Optional[ExtractedText]:
-    try:
-        return await field.get_extracted_text()
-    except Exception:
-        logger.warning(
-            "Error getting extracted text for field. Retrying",
-            exc_info=True,
-            extra={
-                "kbid": field.kbid,
-                "resource_id": field.resource.uuid,
-                "field": f"{field.type}/{field.id}",
-            },
-        )
-        raise
-
-
 async def get_extracted_text_from_field_id(kbid: str, field: FieldId) -> Optional[ExtractedText]:
     rid = field.rid
     orm_resource = await get_resource(kbid, rid)

nucliadb/search/search/{hydrator.py → hydrator/__init__.py}

@@ -28,7 +28,8 @@ from nucliadb.common.external_index_providers.base import TextBlockMatch
 from nucliadb.common.ids import FieldId
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.serialize import managed_serialize
-from nucliadb.search.search import cache, paragraphs
+from nucliadb.search.search import cache
+from nucliadb.search.search.paragraphs import get_paragraph_text
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import ExtractedDataTypeName, Resource
 from nucliadb_models.search import (
@@ -170,7 +171,7 @@ async def hydrate_text_block(
         if concurrency_control is not None:
             await stack.enter_async_context(concurrency_control)
 
-        text_block.text = await paragraphs.get_paragraph_text(
+        text_block.text = await get_paragraph_text(
             kbid=kbid,
             paragraph_id=text_block.paragraph_id,
             highlight=options.highlight,