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

nucliadb/common/datamanagers/synonyms.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
 
 from nucliadb.common.datamanagers.utils import get_kv_pb
 from nucliadb.common.maindb.driver import Transaction
@@ -27,7 +26,7 @@ from nucliadb_protos import knowledgebox_pb2
 KB_SYNONYMS = "/kbs/{kbid}/synonyms"
 
 
-async def get(txn: Transaction, *, kbid: str) -> Optional[knowledgebox_pb2.Synonyms]:
+async def get(txn: Transaction, *, kbid: str) -> knowledgebox_pb2.Synonyms | None:
     key = KB_SYNONYMS.format(kbid=kbid)
     return await get_kv_pb(txn, key, knowledgebox_pb2.Synonyms, for_update=False)
 

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,62 +17,33 @@
 # 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
 
 from nucliadb.common import datamanagers
 from nucliadb.common.external_index_providers.base import ExternalIndexManager
-from nucliadb.common.external_index_providers.pinecone import PineconeIndexManager
-from nucliadb.common.external_index_providers.settings import settings
 from nucliadb_protos.knowledgebox_pb2 import (
-    ExternalIndexProviderType,
     StoredExternalIndexProviderMetadata,
 )
-from nucliadb_utils.utilities import get_endecryptor
 
 
 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).
     """
-    metadata = await get_external_index_metadata(kbid)
-    if metadata is None or metadata.type != ExternalIndexProviderType.PINECONE:
-        # Only Pinecone is supported for now
-        return None
-
-    api_key = get_endecryptor().decrypt(metadata.pinecone_config.encrypted_api_key)
-    default_vectorset = await get_default_vectorset_id(kbid)
-
-    rollover_indexes = None
-    if for_rollover:
-        rollover_metadata = await get_rollover_external_index_metadata(kbid)
-        if rollover_metadata is not None:
-            rollover_indexes = dict(rollover_metadata.pinecone_config.indexes)
-
-    return PineconeIndexManager(
-        kbid=kbid,
-        api_key=api_key,
-        indexes=dict(metadata.pinecone_config.indexes),
-        upsert_parallelism=settings.pinecone_upsert_parallelism,
-        delete_parallelism=settings.pinecone_delete_parallelism,
-        upsert_timeout=settings.pinecone_upsert_timeout,
-        delete_timeout=settings.pinecone_delete_timeout,
-        default_vectorset=default_vectorset,
-        rollover_indexes=rollover_indexes,
-    )
+    return None
 
 
 @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
@@ -96,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/external_index_providers/settings.py

@@ -17,36 +17,10 @@
 # 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 Field
 from pydantic_settings import BaseSettings
 
 
-class ExternalIndexProvidersSettings(BaseSettings):
-    pinecone_upsert_parallelism: int = Field(
-        default=3,
-        title="Pinecone upsert parallelism",
-        description="Number of parallel upserts to Pinecone on each set resource operation",
-    )
-    pinecone_delete_parallelism: int = Field(
-        default=2,
-        title="Pinecone delete parallelism",
-        description="Number of parallel deletes to Pinecone on each delete resource operation",
-    )
-    pinecone_upsert_timeout: float = Field(
-        default=10.0,
-        title="Pinecone upsert timeout",
-        description="Timeout in seconds for each upsert operation to Pinecone",
-    )
-    pinecone_delete_timeout: float = Field(
-        default=10.0,
-        title="Pinecone delete timeout",
-        description="Timeout in seconds for each delete operation to Pinecone",
-    )
-    pinecone_query_timeout: float = Field(
-        default=10.0,
-        title="Pinecone query timeout",
-        description="Timeout in seconds for each query operation to Pinecone",
-    )
+class ExternalIndexProvidersSettings(BaseSettings): ...
 
 
 settings = ExternalIndexProvidersSettings()

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()
@@ -110,7 +95,7 @@ async def parse_expression(
             if rid is None:
                 raise InvalidQueryError("slug", f"Cannot find slug {expr.slug}")
             f.resource.resource_id = rid
-        else:  # pragma: nocover
+        else:  # pragma: no cover
             # Cannot happen due to model validation
             raise ValueError("Resource needs id or slug")
     elif isinstance(expr, Field):
@@ -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/exceptions.py

@@ -21,6 +21,10 @@ class ClientException(Exception):
     pass
 
 
+class ServerException(Exception):
+    pass
+
+
 class NotFoundException(ClientException):
     pass
 
@@ -35,3 +39,7 @@ class RateLimitException(ClientException):
 
 class AccountLimitException(ClientException):
     pass
+
+
+class ServiceUnavailableException(ServerException):
+    pass

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.",
@@ -209,6 +198,10 @@ class ProcessingHTTPClient:
     async def close(self):
         await self.session.close()
 
+    async def reset_session(self):
+        await self.close()
+        self.session = aiohttp.ClientSession()
+
     async def in_progress(self, ack_token: str):
         url = self.base_url_v2 + "/pull/in_progress"
         request = InProgressRequest(ack=[ack_token])
@@ -220,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(
@@ -244,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"
@@ -263,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/http_clients/utils.py

@@ -33,5 +33,8 @@ def check_status(resp: aiohttp.ClientResponse, resp_text: str) -> None:
         raise exceptions.AuthorizationException(f"Unauthorized to access: {resp.status}")
     elif resp.status == 429:
         raise exceptions.RateLimitException("Rate limited")
+    elif resp.status in (502, 503):
+        # Service unavailable, can be retried
+        raise exceptions.ServiceUnavailableException(f"Service unavailable: {resp.status} - {resp_text}")
     else:
         raise exceptions.ClientException(f"Unknown error: {resp.status} - {resp_text}")