nucliadb-models 6.8.1.post4983__py3-none-any.whl → 6.10.0.post5694__py3-none-any.whl

nucliadb_models/hydration.py

@@ -0,0 +1,423 @@
+# Copyright 2025 Bosutech XXI S.L.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from typing import Annotated
+
+from pydantic import BaseModel, Field, StringConstraints
+
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.metadata import Origin
+from nucliadb_models.resource import FieldConversation, FieldFile, FieldLink, FieldText
+from nucliadb_models.search import Image
+from nucliadb_models.security import ResourceSecurity
+
+
+class ResourceHydration(BaseModel, extra="forbid"):
+    title: bool = Field(
+        default=True,
+        description="Hydrate resource titles",
+    )
+    summary: bool = Field(
+        default=False,
+        description="Hydrate resource summaries",
+    )
+
+    origin: bool = Field(
+        default=False,
+        description="Hydrate resource origin",
+    )
+
+    security: bool = Field(
+        default=False,
+        description="Hydrate resource security metadata",
+    )
+
+
+class TextFieldHydration(BaseModel, extra="forbid"):
+    value: bool = Field(
+        default=False,
+        description="Hydrate text field values. Field values are similar payloads to the ones used to create them",
+    )
+    extracted_text: bool = Field(
+        default=False,
+        description="Hydrate extracted text for text fields",
+    )
+    # TODO: what else should be interesting to add?
+
+
+class FileFieldHydration(BaseModel, extra="forbid"):
+    value: bool = Field(
+        default=False,
+        description="Hydrate file field values. Field values are similar payloads to the ones used to create them",
+    )
+    extracted_text: bool = Field(
+        default=False,
+        description="Hydrate extracted text for file fields",
+    )
+    # TODO: what else should be interesting to add?
+
+
+class LinkFieldHydration(BaseModel, extra="forbid"):
+    value: bool = Field(
+        default=False,
+        description="Hydrate link field values. Field values are similar payloads to the ones used to create them",
+    )
+    extracted_text: bool = Field(
+        default=False,
+        description="Hydrate extracted text for link fields",
+    )
+    # TODO: what else should be interesting to add?
+
+
+class ConversationFieldHydration(BaseModel, extra="forbid"):
+    value: bool = Field(
+        default=False,
+        description="Hydrate conversation field values. Field values are similar payloads to the ones used to create them",
+    )
+
+    # TODO: add fields to hydrate conversation fields. Think about how to handle
+    # splits and fulfill the conversational RAG strategies
+
+    # TODO: what else should be interesting to add?
+
+
+class GenericFieldHydration(BaseModel, extra="forbid"):
+    value: bool = Field(
+        default=False,
+        description="Hydrate generic field values. Field values are similar payloads to the ones used to create them",
+    )
+    extracted_text: bool = Field(
+        default=False,
+        description="Hydrate extracted text for generic fields",
+    )
+    # TODO: what else should be interesting to add?
+
+
+class FieldHydration(BaseModel, extra="forbid"):
+    text: TextFieldHydration | None = Field(
+        default_factory=TextFieldHydration,
+        description="Text fields hydration options",
+    )
+    file: FileFieldHydration | None = Field(
+        default_factory=FileFieldHydration,
+        description="File fields hydration options",
+    )
+    link: LinkFieldHydration | None = Field(
+        default_factory=LinkFieldHydration,
+        description="Link fields hydration options",
+    )
+    conversation: ConversationFieldHydration | None = Field(
+        default_factory=ConversationFieldHydration,
+        description="Conversation fields hydration options",
+    )
+    generic: GenericFieldHydration | None = Field(
+        default_factory=GenericFieldHydration,
+        description="Generic fields hydration options",
+    )
+
+
+class NeighbourParagraphHydration(BaseModel, extra="forbid"):
+    before: int = Field(
+        default=2,
+        ge=0,
+        description="Number of previous paragraphs to hydrate",
+    )
+    after: int = Field(
+        default=2,
+        ge=0,
+        description="Number of following paragraphs to hydrate",
+    )
+
+
+class RelatedParagraphHydration(BaseModel, extra="forbid"):
+    neighbours: NeighbourParagraphHydration | None = Field(
+        default=None,
+        description="Hydrate extra paragraphs that surround the original one",
+    )
+
+    # TODO: FEATURE: implement related paragraphs by page
+    # page: bool = Field(
+    #     default=False,
+    #     description="Hydrate all paragraphs in the same page. This only applies to fields with pages",
+    # )
+
+    # TODO: description
+    # XXX: should we let users control the amount of elements?
+    parents: bool = False
+    # TODO: description
+    # XXX: should we let users control the amount of elements?
+    siblings: bool = False
+    # TODO: description
+    # XXX: should we let users control the amount of elements?
+    replacements: bool = False
+
+
+class ImageParagraphHydration(BaseModel, extra="forbid"):
+    # The source image is also known as reference or reference_file in the
+    # paragraph context. The reference/reference_file is the filename of the
+    # source image from which the paragraph has been extracted
+    source_image: bool = Field(
+        default=False,
+        description=(
+            "When a paragraph has been extracted from an image (using OCR, inception...), "
+            "hydrate the image that represents it"
+        ),
+    )
+
+
+class TableParagraphHydration(BaseModel, extra="forbid"):
+    # TODO: implement. ARAG uses the label "/k/table" to check whether a
+    # paragraph is or a table or not. We can also use info on maindb
+    table_page_preview: bool = Field(
+        default=False,
+        description="Hydrate the page preview for the table. This will only hydrate fields with pages",
+    )
+
+
+class ParagraphPageHydration(BaseModel, extra="forbid"):
+    # For some field types (file and link) learning generates previews. A
+    # preview is a PDF file representing the content. For a docx for example, is
+    # the PDF equivalent. Depending on the field type, the preview can
+    # represent, for example, a page in a document or a portion of a webpage.
+    page_with_visual: bool = Field(
+        default=False,
+        description=(
+            "When a paragraph has been extracted from a page containing visual "
+            "content (images, tables...), hydrate the preview of the paragraph's "
+            "page as an image. Not all field types have previews nor visual content"
+        ),
+    )
+
+
+class ParagraphHydration(BaseModel, extra="forbid"):
+    text: bool = Field(
+        default=True,
+        description="Hydrate paragraph text",
+    )
+    image: ImageParagraphHydration | None = Field(
+        default=None,
+        description="Hydrate options for paragraphs extracted from images (using OCR, inception...)",
+    )
+    table: TableParagraphHydration | None = Field(
+        default=None,
+        description="Hydrate options for paragraphs extracted from tables",
+    )
+
+    # TODO: at some point, we should add hydration options for paragraphs from
+    # audio and video
+
+    page: ParagraphPageHydration | None = Field(
+        default=None,
+        description="Hydrte options for paragraphs within a page. This applies to paragraphs in fields with pages",
+    )
+
+    related: RelatedParagraphHydration | None = Field(
+        default=None,
+        description="Hydration options for related paragraphs. For example, neighbours or sibling paragraphs",
+    )
+
+
+class Hydration(BaseModel, extra="forbid"):
+    resource: ResourceHydration | None = Field(
+        default_factory=ResourceHydration,
+        description="Resource hydration options",
+    )
+    field: FieldHydration = Field(
+        default_factory=FieldHydration,
+        description="Field hydration options",
+    )
+    paragraph: ParagraphHydration = Field(
+        default_factory=ParagraphHydration,
+        description="Paragraph hydration options",
+    )
+
+
+ParagraphId = Annotated[
+    str,
+    StringConstraints(
+        # resource-uuid/field-type/field-id/[split-id/]paragraph-id
+        pattern=r"^[0-9a-f]{32}/[acftu]/[a-zA-Z0-9:_-]+(/[^/]{1,128})?/[0-9]+-[0-9]+$",
+        min_length=32 + 1 + 1 + 1 + 1 + 0 + 0 + 1 + 3,
+        # max field id of 250 and 10 digit paragraphs. More than enough
+        max_length=32 + 1 + 1 + 1 + 250 + 1 + 128 + 1 + 21,
+    ),
+]
+
+
+class HydrateRequest(BaseModel, extra="forbid"):
+    data: list[ParagraphId] = Field(
+        description="List of paragraph ids we want to hydrate",
+        max_length=50,
+    )
+    hydration: Hydration = Field(description="Description of how hydration must be performed")
+
+
+### Response models
+
+
+class HydratedResource(BaseModel, extra="forbid"):
+    id: str = Field(description="Unique resource id")
+    slug: str = Field(description="Resource slug")
+
+    title: str | None = None
+    summary: str | None = None
+
+    origin: Origin | None = None
+
+    security: ResourceSecurity | None = None
+
+    # TODO: add resource labels to hydrated resources
+
+
+class FieldExtractedData(BaseModel, extra="forbid"):
+    text: str | None = None
+
+
+class SplitFieldExtractedData(BaseModel, extra="forbid"):
+    texts: dict[str, str] | None = None
+
+
+class HydratedTextField(BaseModel, extra="forbid"):
+    id: str = Field("Unique field id")
+    resource: str = Field("Field resource id")
+    field_type: FieldTypeName = FieldTypeName.TEXT
+
+    value: FieldText | None = None
+    extracted: FieldExtractedData | None = None
+
+
+class HydratedFileField(BaseModel, extra="forbid"):
+    id: str = Field("Unique field id")
+    resource: str = Field("Field resource id")
+    field_type: FieldTypeName = FieldTypeName.FILE
+
+    value: FieldFile | None = None
+    extracted: FieldExtractedData | None = None
+
+    previews: dict[str, Image] | None = Field(
+        default=None,
+        title="Previews of specific parts of the field",
+        description=(
+            "Previews for specific pages of this field. Previews are differents"
+            "depending on the file type. For example, for a PDF file, a preview"
+            "will be an image of a single page."
+            "In this field, previews will be populated according to the hydration"
+            "options requested."
+        ),
+    )
+
+
+class HydratedLinkField(BaseModel, extra="forbid"):
+    id: str = Field("Unique field id")
+    resource: str = Field("Field resource id")
+    field_type: FieldTypeName = FieldTypeName.LINK
+
+    value: FieldLink | None = None
+    extracted: FieldExtractedData | None = None
+
+
+class HydratedConversationField(BaseModel, extra="forbid"):
+    id: str = Field("Unique field id")
+    resource: str = Field("Field resource id")
+    field_type: FieldTypeName = FieldTypeName.CONVERSATION
+
+    value: FieldConversation | None = None
+    extracted: FieldExtractedData | None = None
+
+
+class HydratedGenericField(BaseModel, extra="forbid"):
+    id: str = Field("Unique field id")
+    resource: str = Field("Field resource id")
+    field_type: FieldTypeName = FieldTypeName.TEXT
+
+    value: str | None = None
+    extracted: FieldExtractedData | None = None
+
+
+class RelatedNeighbourParagraphRefs(BaseModel, extra="forbid"):
+    before: list[str] | None = None
+    after: list[str] | None = None
+
+
+class RelatedParagraphRefs(BaseModel, extra="forbid"):
+    neighbours: RelatedNeighbourParagraphRefs | None = None
+    parents: list[str] | None = None
+    siblings: list[str] | None = None
+    replacements: list[str] | None = None
+
+
+class HydratedParagraphImage(BaseModel, extra="forbid"):
+    source_image: Image | None = Field(
+        default=None,
+        description=(
+            "Source image for this paragraph. This only applies to paragraphs "
+            "extracted from an image using OCR or inception, and if this "
+            "hydration option has been enabled in the request"
+        ),
+    )
+
+
+class HydratedParagraphTable(BaseModel, extra="forbid"):
+    page_preview_ref: str | None = Field(
+        default=None,
+        description=(
+            "Referento to the page preview for this paragraph. The actual "
+            "preview will be found in the previews of its field. This only "
+            "applies to paragraphs generated from a table and if the "
+            "corresponding hydration option has been enabled in the request"
+        ),
+    )
+
+
+class HydratedParagraphPage(BaseModel, extra="forbid"):
+    page_preview_ref: str | None = Field(
+        default=None,
+        description=(
+            "Reference to the page preview for this paragraph. The actual "
+            "preview will be found in the previews of its field. This only "
+            "applies to paragraphs extracted from a page containing visual "
+            "content and if the corresponding hydration option has been enabled "
+            "in the request"
+        ),
+    )
+
+
+class HydratedParagraph(BaseModel, extra="forbid"):
+    id: str = Field(description="Unique paragraph id")
+    field: str = Field(description="Paragraph field id")
+    resource: str = Field(description="Paragraph resource id")
+
+    text: str | None = None
+
+    # TODO: add labels to hydrated paragraphs
+    # labels: Optional[list[str]] = None
+
+    related: RelatedParagraphRefs | None = None
+
+    image: HydratedParagraphImage | None = None
+    table: HydratedParagraphTable | None = None
+    page: HydratedParagraphPage | None = None
+
+
+class Hydrated(BaseModel, extra="forbid"):
+    resources: dict[str, HydratedResource]
+    fields: dict[
+        str,
+        HydratedTextField
+        | HydratedFileField
+        | HydratedLinkField
+        | HydratedConversationField
+        | HydratedGenericField,
+    ]
+    paragraphs: dict[str, HydratedParagraph]

nucliadb_models/internal/predict.py

@@ -19,13 +19,11 @@ Models for Predict API v1.
 ATENTION! Keep these models in sync with models on Predict API
 """
 
-from typing import List, Optional
-
 from pydantic import BaseModel, Field
 
 
 class SentenceSearch(BaseModel):
-    vectors: dict[str, List[float]] = Field(
+    vectors: dict[str, list[float]] = Field(
         default_factory=dict,
         description="Sentence vectors for each semantic model",
         min_length=1,
@@ -45,14 +43,14 @@ class Ner(BaseModel):
 
 
 class TokenSearch(BaseModel):
-    tokens: List[Ner] = []
+    tokens: list[Ner] = []
     time: float
     input_tokens: int = 0
 
 
 class QueryInfo(BaseModel):
-    language: Optional[str]
-    stop_words: List[str] = Field(default_factory=list)
+    language: str | None
+    stop_words: list[str] = Field(default_factory=list)
     semantic_thresholds: dict[str, float] = Field(
         default_factory=dict,
         description="Semantic threshold for each semantic model",
@@ -60,10 +58,10 @@ class QueryInfo(BaseModel):
     )
     visual_llm: bool
     max_context: int
-    entities: Optional[TokenSearch]
-    sentence: Optional[SentenceSearch]
+    entities: TokenSearch | None
+    sentence: SentenceSearch | None
     query: str
-    rephrased_query: Optional[str] = None
+    rephrased_query: str | None = None
 
 
 class RerankModel(BaseModel):

nucliadb_models/internal/shards.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 #
 from enum import Enum
-from typing import List, Optional
 
 from pydantic import BaseModel
 
@@ -58,9 +57,9 @@ class ShardReplica(BaseModel):
 
 class ShardObject(BaseModel):
     shard: str
-    nidx_shard_id: Optional[str]
+    nidx_shard_id: str | None
 
 
 class KnowledgeboxShards(BaseModel):
     kbid: str
-    shards: List[ShardObject]
+    shards: list[ShardObject]

nucliadb_models/labels.py

@@ -14,9 +14,8 @@
 #
 
 from enum import Enum
-from typing import Dict, List, Optional
 
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, Field, model_validator
 from typing_extensions import Self
 
 BASE_LABELS: dict[str, set[str]] = {
@@ -96,18 +95,26 @@ class LabelSetKind(str, Enum):
 
 
 class Label(BaseModel):
-    title: str
-    related: Optional[str] = None
-    text: Optional[str] = None
-    uri: Optional[str] = None
+    title: str = Field(
+        description="Title of the label. This is the display name for the label shown in the UI and also used for searching."
+    )
+    related: str | None = None
+    text: str | None = None
+    uri: str | None = None
 
 
 class LabelSet(BaseModel):
-    title: Optional[str] = None
-    color: Optional[str] = "blue"
+    title: str | None = Field(
+        default=None,
+        description="Title of the labelset. It is a prettier display name for the labelset shown in the UI but it is not intended to be used for searching.",
+    )
+    color: str | None = "blue"
     multiple: bool = True
-    kind: List[LabelSetKind] = []
-    labels: List[Label] = []
+    kind: list[LabelSetKind] = []
+    labels: list[Label] = Field(
+        default_factory=list,
+        description="List of labels in the labelset. The titles of the labels must be unique within the labelset.",
+    )
 
     @model_validator(mode="after")
     def check_unique_labels(self) -> Self:
@@ -123,4 +130,4 @@ class LabelSet(BaseModel):
 
 class KnowledgeBoxLabels(BaseModel):
     uuid: str
-    labelsets: Dict[str, LabelSet] = {}
+    labelsets: dict[str, LabelSet] = {}

nucliadb_models/link.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 #
 from datetime import datetime
-from typing import Dict, Optional
 
 from pydantic import BaseModel, Field
 
@@ -25,19 +24,19 @@ from pydantic import BaseModel, Field
 
 
 class FieldLink(BaseModel):
-    added: Optional[datetime] = None
-    headers: Optional[Dict[str, str]] = None
-    cookies: Optional[Dict[str, str]] = None
-    uri: Optional[str] = None
-    language: Optional[str] = None
-    localstorage: Optional[Dict[str, str]] = None
-    css_selector: Optional[str] = None
-    xpath: Optional[str] = None
-    extract_strategy: Optional[str] = Field(
+    added: datetime | None = None
+    headers: dict[str, str] | None = None
+    cookies: dict[str, str] | None = None
+    uri: str | None = None
+    language: str | None = None
+    localstorage: dict[str, str] | None = None
+    css_selector: str | None = None
+    xpath: str | None = None
+    extract_strategy: str | None = Field(
         default=None,
         description="Id of the Nuclia extract strategy used at processing time. If not set, the default strategy was used. Extract strategies are defined at the learning configuration api.",
     )
-    split_strategy: Optional[str] = Field(
+    split_strategy: str | None = Field(
         default=None,
         description="Id of the Nuclia split strategy used at processing time. If not set, the default strategy was used. Split strategies are defined at the learning configuration api.",
     )
@@ -47,18 +46,18 @@ class FieldLink(BaseModel):
 
 
 class LinkField(BaseModel):
-    headers: Optional[Dict[str, str]] = {}
-    cookies: Optional[Dict[str, str]] = {}
+    headers: dict[str, str] | None = {}
+    cookies: dict[str, str] | None = {}
     uri: str
-    language: Optional[str] = None
-    localstorage: Optional[Dict[str, str]] = {}
-    css_selector: Optional[str] = None
-    xpath: Optional[str] = None
-    extract_strategy: Optional[str] = Field(
+    language: str | None = None
+    localstorage: dict[str, str] | None = {}
+    css_selector: str | None = None
+    xpath: str | None = None
+    extract_strategy: str | None = Field(
         default=None,
         description="Id of the Nuclia extract strategy to use at processing time. If not set, the default strategy will be used. Extract strategies are defined at the learning configuration api.",
     )
-    split_strategy: Optional[str] = Field(
+    split_strategy: str | None = Field(
         default=None,
         description="Id of the Nuclia split strategy used at processing time. If not set, the default strategy was used. Split strategies are defined at the learning configuration api.",
     )