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

nucliadb/common/datamanagers/utils.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import contextlib
-from typing import Optional, Type, TypeVar
+from typing import TypeVar
 
 from google.protobuf.message import Message
 
@@ -29,9 +29,9 @@ PB_TYPE = TypeVar("PB_TYPE", bound=Message)
 
 
 async def get_kv_pb(
-    txn: Transaction, key: str, pb_type: Type[PB_TYPE], for_update: bool = True
-) -> Optional[PB_TYPE]:
-    serialized: Optional[bytes] = await txn.get(key, for_update=for_update)
+    txn: Transaction, key: str, pb_type: type[PB_TYPE], for_update: bool = True
+) -> PB_TYPE | None:
+    serialized: bytes | None = await txn.get(key, for_update=for_update)
     if serialized is None:
         return None
     pb = pb_type()

nucliadb/common/datamanagers/vectorsets.py

@@ -17,7 +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 typing import AsyncIterator, Optional
+from collections.abc import AsyncIterator
 
 from nucliadb.common.datamanagers.utils import get_kv_pb
 from nucliadb.common.maindb.driver import Transaction
@@ -37,7 +37,7 @@ async def initialize(txn: Transaction, *, kbid: str):
 
 async def get(
     txn: Transaction, *, kbid: str, vectorset_id: str
-) -> Optional[knowledgebox_pb2.VectorSetConfig]:
+) -> knowledgebox_pb2.VectorSetConfig | None:
     kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
     index = _find_vectorset(kb_vectorsets, vectorset_id)
     if index is None:
@@ -80,7 +80,7 @@ async def set(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.VectorSet
 
 async def delete(
     txn: Transaction, *, kbid: str, vectorset_id: str
-) -> Optional[knowledgebox_pb2.VectorSetConfig]:
+) -> knowledgebox_pb2.VectorSetConfig | None:
     kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
     index = _find_vectorset(kb_vectorsets, vectorset_id)
     if index is None:
@@ -111,7 +111,7 @@ async def _get_or_default(
 
 def _find_vectorset(
     kb_vectorsets: knowledgebox_pb2.KnowledgeBoxVectorSetsConfig, vectorset_id: str
-) -> Optional[int]:
+) -> int | None:
     """Return the position of the vectorset in `vectorsets` or `None` if not found."""
     for idx, vectorset in enumerate(kb_vectorsets.vectorsets):
         if vectorset.vectorset_id == vectorset_id:

nucliadb/common/external_index_providers/base.py

@@ -19,8 +19,9 @@
 #
 import abc
 import logging
+from collections.abc import Iterator
 from dataclasses import dataclass
-from typing import Any, Iterator, Optional
+from typing import Any
 
 from nidx_protos.nodereader_pb2 import SearchRequest
 from nidx_protos.noderesources_pb2 import Resource
@@ -30,7 +31,9 @@ from nucliadb.common.counters import IndexCounts
 from nucliadb.common.external_index_providers.exceptions import ExternalIndexingError
 from nucliadb.common.ids import ParagraphId
 from nucliadb_models.external_index_providers import ExternalIndexProviderType
+from nucliadb_models.retrieval import Score
 from nucliadb_models.search import SCORE_TYPE, Relations, TextPosition
+from nucliadb_protos import resources_pb2
 from nucliadb_protos.knowledgebox_pb2 import (
     CreateExternalIndexProviderMetadata,
     StoredExternalIndexProviderMetadata,
@@ -43,6 +46,16 @@ logger = logging.getLogger(__name__)
 manager_observer = Observer("external_index_manager", labels={"operation": "", "provider": ""})
 
 
+# /k/ocr
+_OCR_LABEL = (
+    f"/k/{resources_pb2.Paragraph.TypeParagraph.Name(resources_pb2.Paragraph.TypeParagraph.OCR).lower()}"
+)
+# /k/inception
+_INCEPTION_LABEL = (
+    f"/k/{resources_pb2.Paragraph.TypeParagraph.Name(resources_pb2.Paragraph.TypeParagraph.OCR).lower()}"
+)
+
+
 @dataclass
 class VectorsetExternalIndex:
     """
@@ -57,9 +70,19 @@ class VectorsetExternalIndex:
 
 class ScoredTextBlock(BaseModel):
     paragraph_id: ParagraphId
-    score: float
     score_type: SCORE_TYPE
 
+    scores: list[Score]
+
+    @property
+    def score(self) -> float:
+        return self.current_score.score
+
+    @property
+    def current_score(self) -> Score:
+        assert len(self.scores) > 0, "text block matches must be scored"
+        return self.scores[-1]
+
 
 class TextBlockMatch(ScoredTextBlock):
     """
@@ -72,11 +95,15 @@ class TextBlockMatch(ScoredTextBlock):
     page_with_visual: bool = False
     fuzzy_search: bool
     is_a_table: bool = False
-    representation_file: Optional[str] = None
+    representation_file: str | None = None
     paragraph_labels: list[str] = []
     field_labels: list[str] = []
-    text: Optional[str] = None
-    relevant_relations: Optional[Relations] = None
+    text: str | None = None
+    relevant_relations: Relations | None = None
+
+    @property
+    def is_an_image(self) -> bool:
+        return _OCR_LABEL in self.paragraph_labels or _INCEPTION_LABEL in self.paragraph_labels
 
 
 class QueryResults(BaseModel):

nucliadb/common/external_index_providers/manager.py

@@ -17,7 +17,6 @@
 # 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 Optional
 
 import async_lru
 
@@ -30,7 +29,7 @@ from nucliadb_protos.knowledgebox_pb2 import (
 
 async def get_external_index_manager(
     kbid: str, for_rollover: bool = False
-) -> Optional[ExternalIndexManager]:
+) -> ExternalIndexManager | None:
     """
     Returns an ExternalIndexManager for the given kbid.
     If for_rollover is True, the ExternalIndexManager returned will include the rollover indexes (if any).
@@ -39,12 +38,12 @@ async def get_external_index_manager(
 
 
 @async_lru.alru_cache(maxsize=None)
-async def get_external_index_metadata(kbid: str) -> Optional[StoredExternalIndexProviderMetadata]:
+async def get_external_index_metadata(kbid: str) -> StoredExternalIndexProviderMetadata | None:
     return await datamanagers.atomic.kb.get_external_index_provider_metadata(kbid=kbid)
 
 
 @async_lru.alru_cache(maxsize=None)
-async def get_default_vectorset_id(kbid: str) -> Optional[str]:
+async def get_default_vectorset_id(kbid: str) -> str | None:
     """
     While we are transitioning to the new vectorset system, we need to take into account
     that KBs that have only one semantic model will have the `vectorset_id` field on BrokerMessage.field_vectors
@@ -68,6 +67,6 @@ async def get_default_vectorset_id(kbid: str) -> Optional[str]:
 
 async def get_rollover_external_index_metadata(
     kbid: str,
-) -> Optional[StoredExternalIndexProviderMetadata]:
+) -> StoredExternalIndexProviderMetadata | None:
     async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.rollover.get_kb_rollover_external_index_metadata(txn, kbid=kbid)

nucliadb/common/filter_expression.py

@@ -18,13 +18,14 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Union
 
 from nidx_protos.nodereader_pb2 import FilterExpression as PBFilterExpression
+from typing_extensions import assert_never
 
 from nucliadb.common import datamanagers
 from nucliadb.common.exceptions import InvalidQueryError
 from nucliadb.common.ids import FIELD_TYPE_NAME_TO_STR
+from nucliadb_models.common import Paragraph
 from nucliadb_models.filters import (
     And,
     DateCreated,
@@ -50,44 +51,28 @@ from nucliadb_models.filters import (
     ResourceMimetype,
     Status,
 )
+from nucliadb_models.metadata import ResourceProcessingStatus
 
 # Filters that end up as a facet
-FacetFilter = Union[
-    OriginTag,
-    Label,
-    ResourceMimetype,
-    FieldMimetype,
-    Entity,
-    Language,
-    OriginMetadata,
-    OriginPath,
-    Generated,
-    Kind,
-    OriginCollaborator,
-    OriginSource,
-    Status,
-]
-# In Python 3.9 we cannot do isinstance against an union
-# Once we support only 3.10+, we can remove this
-FacetFilterTypes = (
-    OriginTag,
-    Label,
-    ResourceMimetype,
-    FieldMimetype,
-    Entity,
-    Language,
-    OriginMetadata,
-    OriginPath,
-    Generated,
-    Kind,
-    OriginCollaborator,
-    OriginSource,
-    Status,
+FacetFilter = (
+    OriginTag
+    | Label
+    | ResourceMimetype
+    | FieldMimetype
+    | Entity
+    | Language
+    | OriginMetadata
+    | OriginPath
+    | Generated
+    | Kind
+    | OriginCollaborator
+    | OriginSource
+    | Status
 )
 
 
 async def parse_expression(
-    expr: Union[FieldFilterExpression, ParagraphFilterExpression],
+    expr: FieldFilterExpression | ParagraphFilterExpression,
     kbid: str,
 ) -> PBFilterExpression:
     f = PBFilterExpression()
@@ -131,12 +116,10 @@ async def parse_expression(
             f.date.since.FromDatetime(expr.since)
         if expr.until:
             f.date.until.FromDatetime(expr.until)
-    elif isinstance(expr, FacetFilterTypes):
+    elif isinstance(expr, FacetFilter):
         f.facet.facet = facet_from_filter(expr)
     else:
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+        assert_never(expr)
 
     return f
 
@@ -190,13 +173,118 @@ def facet_from_filter(expr: FacetFilter) -> str:
     elif isinstance(expr, Status):
         facet = f"/n/s/{expr.status.value}"
     else:
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+        assert_never(expr)
 
     return facet
 
 
+def filter_from_facet(facet: str) -> FacetFilter:
+    expr: FacetFilter
+
+    if facet.startswith("/t/"):
+        value = facet.removeprefix("/t/")
+        expr = OriginTag(tag=value)
+
+    elif facet.startswith("/l/"):
+        value = facet.removeprefix("/l/")
+        parts = value.split("/", maxsplit=1)
+        if len(parts) == 1:
+            type = parts[0]
+            expr = Label(labelset=type)
+        else:
+            type, subtype = parts
+            expr = Label(labelset=type, label=subtype)
+
+    elif facet.startswith("/n/i/"):
+        value = facet.removeprefix("/n/i/")
+        parts = value.split("/", maxsplit=1)
+        if len(parts) == 1:
+            type = parts[0]
+            expr = ResourceMimetype(type=type)
+        else:
+            type, subtype = parts
+            expr = ResourceMimetype(type=type, subtype=subtype)
+
+    elif facet.startswith("/mt/"):
+        value = facet.removeprefix("/mt/")
+        parts = value.split("/", maxsplit=1)
+        if len(parts) == 1:
+            type = parts[0]
+            expr = FieldMimetype(type=type)
+        else:
+            type, subtype = parts
+            expr = FieldMimetype(type=type, subtype=subtype)
+
+    elif facet.startswith("/e/"):
+        value = facet.removeprefix("/e/")
+        parts = value.split("/", maxsplit=1)
+        if len(parts) == 1:
+            subtype = parts[0]
+            expr = Entity(subtype=subtype)
+        else:
+            subtype, value = parts
+            expr = Entity(subtype=subtype, value=value)
+
+    elif facet.startswith("/s/p"):
+        value = facet.removeprefix("/s/p/")
+        expr = Language(language=value, only_primary=True)
+
+    elif facet.startswith("/s/s"):
+        value = facet.removeprefix("/s/s/")
+        expr = Language(language=value, only_primary=False)
+
+    elif facet.startswith("/m/"):
+        value = facet.removeprefix("/m/")
+        parts = value.split("/", maxsplit=1)
+        if len(parts) == 1:
+            field = parts[0]
+            expr = OriginMetadata(field=field)
+        else:
+            field, value = parts
+            expr = OriginMetadata(field=field, value=value)
+
+    elif facet.startswith("/p/"):
+        value = facet.removeprefix("/p/")
+        expr = OriginPath(prefix=value)
+
+    elif facet.startswith("/g/da"):
+        value = facet.removeprefix("/g/da")
+        expr = expr = Generated(by="data-augmentation")
+        if value.removeprefix("/"):
+            expr.da_task = value.removeprefix("/")
+
+    elif facet.startswith("/k/"):
+        value = facet.removeprefix("/k/")
+        try:
+            kind = Paragraph.TypeParagraph(value.upper())
+        except ValueError:
+            raise InvalidQueryError("filters", f"invalid paragraph kind: {value}")
+        expr = Kind(kind=kind)
+
+    elif facet.startswith("/u/o/"):
+        value = facet.removeprefix("/u/o/")
+        expr = OriginCollaborator(collaborator=value)
+
+    elif facet.startswith("/u/s"):
+        value = facet.removeprefix("/u/s")
+        expr = OriginSource()
+        if value.removeprefix("/"):
+            expr.id = value.removeprefix("/")
+
+    elif facet.startswith("/n/s/"):
+        value = facet.removeprefix("/n/s/")
+        try:
+            status = ResourceProcessingStatus(value.upper())
+        except ValueError:
+            raise InvalidQueryError("filters", f"invalid resource processing status: {value}")
+        expr = Status(status=status)
+
+    else:
+        raise InvalidQueryError("filters", f"invalid filter: {facet}")
+
+    return expr
+
+
 def add_and_expression(dest: PBFilterExpression, add: PBFilterExpression):
     dest_expr_type = dest.WhichOneof("expr")
     if dest_expr_type is None:

nucliadb/common/http_clients/processing.py

@@ -19,10 +19,8 @@
 #
 import logging
 from datetime import datetime
-from typing import Optional
 
 import aiohttp
-import jwt
 import pydantic
 
 from nucliadb_utils.helpers import MessageProgressUpdater
@@ -33,15 +31,6 @@ from .utils import check_status
 logger = logging.getLogger(__name__)
 
 
-def get_nua_api_id() -> str:
-    assert nuclia_settings.nuclia_service_account is not None
-    claimset = jwt.decode(
-        nuclia_settings.nuclia_service_account,
-        options={"verify_signature": False},
-    )
-    return claimset.get("sub")
-
-
 def get_processing_api_url() -> str:
     if nuclia_settings.nuclia_service_account:
         return (
@@ -64,10 +53,10 @@ def get_processing_api_v2_url() -> str:
 
 class PullResponse(pydantic.BaseModel):
     status: str
-    payload: Optional[str] = None
+    payload: str | None = None
     payloads: list[bytes] = []
-    msgid: Optional[str] = None
-    cursor: Optional[int] = None
+    msgid: str | None = None
+    cursor: int | None = None
 
 
 class PullPosition(pydantic.BaseModel):
@@ -86,7 +75,7 @@ class RequestsResult(pydantic.BaseModel):
         description="Resource ID.",
     )
     kbid: str = pydantic.Field(..., title="KnowledgeBox ID")
-    title: Optional[str] = pydantic.Field(
+    title: str | None = pydantic.Field(
         None,
         title="Title",
         description="Title of the resource.",
@@ -111,12 +100,12 @@ class RequestsResult(pydantic.BaseModel):
         title="Timestamp",
         description="Timestamp of when the resource was first scheduled.",
     )
-    completed_at: Optional[datetime] = pydantic.Field(
+    completed_at: datetime | None = pydantic.Field(
         None,
         title="Completed At",
         description="Timestamp of when the resource was completed",
     )
-    scheduled_at: Optional[datetime] = pydantic.Field(
+    scheduled_at: datetime | None = pydantic.Field(
         None,
         title="Scheduled At",
         description="Timestamp of when the resource was first scheduled.",
@@ -149,7 +138,7 @@ class RequestsResults(pydantic.BaseModel):
         title="Results",
         description="List of results.",
     )
-    cursor: Optional[str] = pydantic.Field(
+    cursor: str | None = pydantic.Field(
         None,
         title="Cursor",
         description="Cursor to use for the next page of results.",
@@ -224,7 +213,7 @@ class ProcessingHTTPClient:
 
     async def pull_v2(
         self, ack_tokens: list[str], limit: int = 1, timeout: float = 5
-    ) -> Optional[PullResponseV2]:
+    ) -> PullResponseV2 | None:
         url = self.base_url_v2 + "/pull"
         request = PullRequestV2(limit=limit, timeout=timeout, ack=ack_tokens)
         async with self.session.post(
@@ -248,9 +237,9 @@ class ProcessingHTTPClient:
 
     async def requests(
         self,
-        cursor: Optional[str] = None,
-        scheduled: Optional[bool] = None,
-        kbid: Optional[str] = None,
+        cursor: str | None = None,
+        scheduled: bool | None = None,
+        kbid: str | None = None,
         limit: int = 20,
     ) -> RequestsResults:
         url = self.base_url + "/requests"
@@ -267,7 +256,7 @@ class ProcessingHTTPClient:
             check_status(resp, resp_text)
             return RequestsResults.model_validate_json(resp_text)
 
-    async def stats(self, kbid: str, timeout: Optional[float] = 1.0) -> StatsResponse:
+    async def stats(self, kbid: str, timeout: float | None = 1.0) -> StatsResponse:
         url = self.base_url + "/stats"
         async with self.session.get(
             url,

nucliadb/common/ids.py

@@ -24,7 +24,6 @@ paragraphs... Avoiding spread of id construction and parsing everywhere
 """
 
 from dataclasses import dataclass
-from typing import Optional
 
 from nucliadb_models.common import FieldTypeName
 from nucliadb_protos.resources_pb2 import FieldType
@@ -77,7 +76,7 @@ class FieldId:
     type: str
     key: str
     # also knwon as `split`, this indicates a part of a field in, for example, conversations
-    subfield_id: Optional[str] = None
+    subfield_id: str | None = None
 
     @classmethod
     def from_string(cls, value: str) -> "FieldId":
@@ -113,7 +112,7 @@ class FieldId:
 
     @classmethod
     def from_pb(
-        cls, rid: str, field_type: FieldType.ValueType, key: str, subfield_id: Optional[str] = None
+        cls, rid: str, field_type: FieldType.ValueType, key: str, subfield_id: str | None = None
     ) -> "FieldId":
         return cls(rid=rid, type=FIELD_TYPE_PB_TO_STR[field_type], key=key, subfield_id=subfield_id)
 
@@ -127,6 +126,9 @@ class FieldId:
         else:
             return f"{self.rid}/{self.type}/{self.key}/{self.subfield_id}"
 
+    def full_without_subfield(self) -> str:
+        return f"{self.rid}/{self.type}/{self.key}"
+
     def short_without_subfield(self) -> str:
         return f"/{self.type}/{self.key}"
 
@@ -262,7 +264,7 @@ class VectorId:
         return hash(self.full())
 
 
-def extract_data_augmentation_id(generated_field_id: str) -> Optional[str]:
+def extract_data_augmentation_id(generated_field_id: str) -> str | None:
     """Data augmentation generated fields have a strict id with the following
     format:
     `da-{task_id}-{original:field_type}-{original:field_id}[-{original:split}]`

nucliadb/common/locking.py

@@ -22,7 +22,6 @@ import logging
 import time
 import uuid
 from dataclasses import dataclass
-from typing import Optional
 
 import orjson
 
@@ -99,7 +98,7 @@ class _Lock:
         self.task = asyncio.create_task(self._refresh_task())
         return self
 
-    async def get_lock_data(self, txn: Transaction) -> Optional[LockValue]:
+    async def get_lock_data(self, txn: Transaction) -> LockValue | None:
         existing_data = await txn.get(self.key, for_update=True)
         if existing_data is None:
             return None

nucliadb/common/maindb/driver.py

@@ -20,8 +20,9 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Optional
+from typing import ClassVar
 
 DEFAULT_SCAN_LIMIT = -1
 DEFAULT_BATCH_SCAN_LIMIT = 500
@@ -37,10 +38,10 @@ class Transaction:
     async def commit(self):
         raise NotImplementedError()
 
-    async def batch_get(self, keys: list[str], for_update: bool = False) -> list[Optional[bytes]]:
+    async def batch_get(self, keys: list[str], for_update: bool = False) -> list[bytes | None]:
         raise NotImplementedError()
 
-    async def get(self, key: str, for_update: bool = False) -> Optional[bytes]:
+    async def get(self, key: str, for_update: bool = False) -> bytes | None:
         raise NotImplementedError()
 
     async def set(self, key: str, value: bytes):
@@ -57,7 +58,7 @@ class Transaction:
 
     def keys(
         self, match: str, count: int = DEFAULT_SCAN_LIMIT, include_start: bool = True
-    ) -> AsyncGenerator[str, None]:
+    ) -> AsyncGenerator[str]:
         raise NotImplementedError()
 
     async def count(self, match: str) -> int:
@@ -66,7 +67,7 @@ class Transaction:
 
 class Driver:
     initialized = False
-    _abort_tasks: list[asyncio.Task] = []
+    _abort_tasks: ClassVar[list[asyncio.Task]] = []
 
     async def initialize(self):
         raise NotImplementedError()
@@ -81,15 +82,15 @@ class Driver:
                     pass
 
     @asynccontextmanager
-    async def _transaction(self, *, read_only: bool) -> AsyncGenerator[Transaction, None]:
+    async def _transaction(self, *, read_only: bool) -> AsyncGenerator[Transaction]:
         yield Transaction()
 
     @asynccontextmanager
-    async def ro_transaction(self) -> AsyncGenerator[Transaction, None]:
+    async def ro_transaction(self) -> AsyncGenerator[Transaction]:
         async with self._transaction(read_only=True) as txn:
             yield txn
 
     @asynccontextmanager
-    async def rw_transaction(self) -> AsyncGenerator[Transaction, None]:
+    async def rw_transaction(self) -> AsyncGenerator[Transaction]:
         async with self._transaction(read_only=False) as txn:
             yield txn

nucliadb/common/maindb/local.py

@@ -19,8 +19,8 @@
 #
 import glob
 import os
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Optional
 
 from nucliadb.common.maindb.driver import (
     DEFAULT_BATCH_SCAN_LIMIT,
@@ -78,7 +78,7 @@ class LocalTransaction(Transaction):
             # Deleting a key that does not exist
             pass
 
-    async def read(self, key: str) -> Optional[bytes]:
+    async def read(self, key: str) -> bytes | None:
         try:
             async with aiofiles.open(self.compute_path(key), "rb") as resp:
                 return await resp.read()
@@ -106,8 +106,8 @@ class LocalTransaction(Transaction):
         self.clean()
         self.open = False
 
-    async def batch_get(self, keys: list[str], for_update: bool = False) -> list[Optional[bytes]]:
-        results: list[Optional[bytes]] = []
+    async def batch_get(self, keys: list[str], for_update: bool = False) -> list[bytes | None]:
+        results: list[bytes | None] = []
         for key in keys:
             obj = await self.get(key)
             if obj:
@@ -125,7 +125,7 @@ class LocalTransaction(Transaction):
 
         return results
 
-    async def get(self, key: str, for_update: bool = False) -> Optional[bytes]:
+    async def get(self, key: str, for_update: bool = False) -> bytes | None:
         if key in self.deleted_keys:
             raise KeyError(f"Not found {key}")