nucliadb 6.3.7.post4116__py3-none-any.whl → 6.3.7.post4120__py3-none-any.whl

nucliadb/search/search/chat/query.py

@@ -20,6 +20,11 @@
 import asyncio
 from typing import Iterable, Optional, Union
 
+from nidx_protos.nodereader_pb2 import (
+    GraphSearchResponse,
+    SearchResponse,
+)
+
 from nucliadb.common.models_utils import to_proto
 from nucliadb.search import logger
 from nucliadb.search.predict import AnswerStatusCode
@@ -52,10 +57,6 @@ from nucliadb_models.search import (
     parse_rephrase_prompt,
 )
 from nucliadb_protos import audit_pb2
-from nucliadb_protos.nodereader_pb2 import (
-    GraphSearchResponse,
-    SearchResponse,
-)
 from nucliadb_protos.utils_pb2 import RelationNode
 from nucliadb_telemetry.errors import capture_exception
 from nucliadb_utils.utilities import get_audit

nucliadb/search/search/fetch.py

@@ -21,6 +21,8 @@ import asyncio
 from contextvars import ContextVar
 from typing import Optional
 
+from nidx_protos.nodereader_pb2 import DocumentResult, ParagraphResult
+
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
@@ -30,7 +32,6 @@ from nucliadb.search.search import cache
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import ExtractedDataTypeName, Resource
 from nucliadb_models.search import ResourceProperties
-from nucliadb_protos.nodereader_pb2 import DocumentResult, ParagraphResult
 from nucliadb_protos.resources_pb2 import Paragraph
 from nucliadb_utils import const
 from nucliadb_utils.utilities import has_feature

nucliadb/search/search/find_merge.py

@@ -20,6 +20,15 @@
 import asyncio
 from typing import Iterable, Optional, Union
 
+from nidx_protos.nodereader_pb2 import (
+    DocumentScored,
+    GraphSearchResponse,
+    ParagraphResult,
+    ParagraphSearchResponse,
+    SearchResponse,
+    VectorSearchResponse,
+)
+
 from nucliadb.common.external_index_providers.base import TextBlockMatch
 from nucliadb.common.ids import ParagraphId, VectorId
 from nucliadb.search import SERVICE_NAME, logger
@@ -50,14 +59,6 @@ from nucliadb_models.search import (
     ResourceProperties,
     TextPosition,
 )
-from nucliadb_protos.nodereader_pb2 import (
-    DocumentScored,
-    GraphSearchResponse,
-    ParagraphResult,
-    ParagraphSearchResponse,
-    SearchResponse,
-    VectorSearchResponse,
-)
 from nucliadb_telemetry import metrics
 
 from .metrics import merge_observer

nucliadb/search/search/graph_merge.py

@@ -19,6 +19,8 @@
 #
 
 
+from nidx_protos import nodereader_pb2
+
 from nucliadb.common.models_utils.from_proto import RelationNodeTypePbMap, RelationTypePbMap
 from nucliadb_models.graph import responses as graph_responses
 from nucliadb_models.graph.responses import (
@@ -26,7 +28,6 @@ from nucliadb_models.graph.responses import (
     GraphRelationsSearchResponse,
     GraphSearchResponse,
 )
-from nucliadb_protos import nodereader_pb2
 
 
 def build_graph_response(results: list[nodereader_pb2.GraphSearchResponse]) -> GraphSearchResponse:

nucliadb/search/search/graph_strategy.py

@@ -21,6 +21,7 @@ import json
 from collections import defaultdict
 from typing import Any, Collection, Iterable, Optional, Union
 
+from nidx_protos import nodereader_pb2
 from nuclia_models.predict.generative_responses import (
     JSONGenerativeResponse,
     MetaGenerativeResponse,
@@ -73,7 +74,6 @@ from nucliadb_models.search import (
     TextPosition,
     UserPrompt,
 )
-from nucliadb_protos import nodereader_pb2
 from nucliadb_protos.utils_pb2 import RelationNode
 
 SCHEMA = {

nucliadb/search/search/merge.py

@@ -22,6 +22,18 @@ import datetime
 import math
 from typing import Any, Iterable, Optional, Set, Union
 
+from nidx_protos.nodereader_pb2 import (
+    DocumentResult,
+    DocumentScored,
+    DocumentSearchResponse,
+    GraphSearchResponse,
+    ParagraphResult,
+    ParagraphSearchResponse,
+    SearchResponse,
+    SuggestResponse,
+    VectorSearchResponse,
+)
+
 from nucliadb.common.ids import FieldId, ParagraphId
 from nucliadb.common.models_utils import from_proto
 from nucliadb.common.models_utils.from_proto import RelationTypePbMap
@@ -61,17 +73,6 @@ from nucliadb_models.search import (
     SortOrder,
     TextPosition,
 )
-from nucliadb_protos.nodereader_pb2 import (
-    DocumentResult,
-    DocumentScored,
-    DocumentSearchResponse,
-    GraphSearchResponse,
-    ParagraphResult,
-    ParagraphSearchResponse,
-    SearchResponse,
-    SuggestResponse,
-    VectorSearchResponse,
-)
 from nucliadb_protos.utils_pb2 import RelationNode
 
 from .metrics import merge_observer

nucliadb/search/search/query.py

@@ -20,6 +20,9 @@
 from datetime import datetime
 from typing import Any, Optional
 
+from nidx_protos import nodereader_pb2
+from nidx_protos.noderesources_pb2 import Resource
+
 from nucliadb.common import datamanagers
 from nucliadb.search.search.filters import (
     translate_label,
@@ -33,8 +36,7 @@ from nucliadb_models.search import (
     SortOrder,
     SuggestOptions,
 )
-from nucliadb_protos import nodereader_pb2, utils_pb2
-from nucliadb_protos.noderesources_pb2 import Resource
+from nucliadb_protos import utils_pb2
 
 from .exceptions import InvalidQueryError
 from .query_parser.filter_expression import add_and_expression, parse_expression

nucliadb/search/search/query_parser/filter_expression.py

@@ -20,6 +20,8 @@
 
 from typing import Union
 
+from nidx_protos.nodereader_pb2 import FilterExpression as PBFilterExpression
+
 from nucliadb.common import datamanagers
 from nucliadb.common.ids import FIELD_TYPE_NAME_TO_STR
 from nucliadb.search.search.exceptions import InvalidQueryError
@@ -48,7 +50,6 @@ from nucliadb_models.filters import (
     ResourceMimetype,
     Status,
 )
-from nucliadb_protos.nodereader_pb2 import FilterExpression as PBFilterExpression
 
 # Filters that end up as a facet
 FacetFilter = Union[

nucliadb/search/search/query_parser/models.py

@@ -21,11 +21,12 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Literal, Optional, Union
 
+from nidx_protos import nodereader_pb2
 from pydantic import BaseModel, ConfigDict, Field
 
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb_models import search as search_models
-from nucliadb_protos import nodereader_pb2, utils_pb2
+from nucliadb_protos import utils_pb2
 
 ### Retrieval
 

nucliadb/search/search/query_parser/old_filters.py

@@ -22,12 +22,13 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Optional, Union
 
+from nidx_protos.nodereader_pb2 import FilterExpression
+
 from nucliadb.search.search.filters import translate_label
 from nucliadb_models.search import (
     Filter,
 )
 from nucliadb_protos import knowledgebox_pb2
-from nucliadb_protos.nodereader_pb2 import FilterExpression
 
 from .exceptions import InvalidQueryError
 from .fetcher import Fetcher

nucliadb/search/search/query_parser/parsers/find.py

@@ -20,6 +20,7 @@
 
 from typing import Optional
 
+from nidx_protos import nodereader_pb2
 from pydantic import ValidationError
 
 from nucliadb.common.models_utils.from_proto import RelationNodeTypeMap
@@ -47,7 +48,7 @@ from nucliadb_models.filters import FilterExpression
 from nucliadb_models.search import (
     FindRequest,
 )
-from nucliadb_protos import nodereader_pb2, utils_pb2
+from nucliadb_protos import utils_pb2
 
 from .common import (
     parse_keyword_query,

nucliadb/search/search/query_parser/parsers/graph.py

@@ -20,13 +20,15 @@
 
 from typing import Optional, Union
 
+from nidx_protos import nodereader_pb2
+
 from nucliadb.common.models_utils.from_proto import RelationNodeTypeMap, RelationTypeMap
 from nucliadb.search.search.query_parser.filter_expression import add_and_expression, parse_expression
 from nucliadb.search.search.query_parser.models import GraphRetrieval
 from nucliadb.search.search.utils import filter_hidden_resources
 from nucliadb_models.graph import requests as graph_requests
 from nucliadb_models.labels import LABEL_HIDDEN
-from nucliadb_protos import nodereader_pb2, utils_pb2
+from nucliadb_protos import utils_pb2
 
 
 async def parse_graph_search(kbid: str, item: graph_requests.GraphSearchRequest) -> GraphRetrieval:

nucliadb/search/search/query_parser/parsers/search.py

@@ -19,6 +19,8 @@
 #
 from typing import Optional
 
+from nidx_protos import nodereader_pb2
+
 from nucliadb.search.search.metrics import query_parser_observer
 from nucliadb.search.search.query import expand_entities
 from nucliadb.search.search.query_parser.exceptions import InvalidQueryError
@@ -42,7 +44,7 @@ from nucliadb_models.search import (
     SortOptions,
     SortOrder,
 )
-from nucliadb_protos import nodereader_pb2, utils_pb2
+from nucliadb_protos import utils_pb2
 
 from .common import (
     parse_keyword_query,

nucliadb/search/search/query_parser/parsers/unit_retrieval.py

@@ -19,6 +19,9 @@
 #
 from typing import Optional
 
+from nidx_protos import nodereader_pb2
+from nidx_protos.nodereader_pb2 import SearchRequest
+
 from nucliadb.search.search.filters import translate_label
 from nucliadb.search.search.metrics import node_features, query_parser_observer
 from nucliadb.search.search.query import apply_entities_filter, get_sort_field_proto
@@ -26,8 +29,7 @@ from nucliadb.search.search.query_parser.filter_expression import add_and_expres
 from nucliadb.search.search.query_parser.models import ParsedQuery, PredictReranker, UnitRetrieval
 from nucliadb_models.labels import LABEL_HIDDEN, translate_system_to_alias_label
 from nucliadb_models.search import SortOrderMap
-from nucliadb_protos import nodereader_pb2, utils_pb2
-from nucliadb_protos.nodereader_pb2 import SearchRequest
+from nucliadb_protos import utils_pb2
 
 
 @query_parser_observer.wrap({"type": "convert_retrieval_to_proto"})

nucliadb/search/search/shards.py

@@ -22,9 +22,7 @@ import asyncio
 import backoff
 from grpc import StatusCode
 from grpc.aio import AioRpcError
-
-from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb_protos.nodereader_pb2 import (
+from nidx_protos.nodereader_pb2 import (
     GetShardRequest,
     GraphSearchRequest,
     GraphSearchResponse,
@@ -33,7 +31,9 @@ from nucliadb_protos.nodereader_pb2 import (
     SuggestRequest,
     SuggestResponse,
 )
-from nucliadb_protos.noderesources_pb2 import Shard
+from nidx_protos.noderesources_pb2 import Shard
+
+from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb_telemetry import metrics
 
 node_observer = metrics.Observer(

nucliadb/train/generators/field_classifier.py

@@ -20,6 +20,8 @@
 
 from typing import AsyncGenerator
 
+from nidx_protos.nodereader_pb2 import StreamRequest
+
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.train import logger
@@ -30,7 +32,6 @@ from nucliadb_protos.dataset_pb2 import (
     TextLabel,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamRequest
 
 
 def field_classification_batch_generator(

nucliadb/train/generators/field_streaming.py

@@ -20,6 +20,8 @@
 
 from typing import AsyncGenerator, Optional
 
+from nidx_protos.nodereader_pb2 import StreamRequest
+
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.train import logger
@@ -29,7 +31,6 @@ from nucliadb_protos.dataset_pb2 import (
     FieldStreamingBatch,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamRequest
 from nucliadb_protos.resources_pb2 import Basic, FieldComputedMetadata
 from nucliadb_protos.utils_pb2 import ExtractedText
 

nucliadb/train/generators/paragraph_classifier.py

@@ -21,6 +21,7 @@
 from typing import AsyncGenerator
 
 from fastapi import HTTPException
+from nidx_protos.nodereader_pb2 import StreamRequest
 
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.train.generators.utils import batchify, get_paragraph
@@ -30,7 +31,6 @@ from nucliadb_protos.dataset_pb2 import (
     TextLabel,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamRequest
 
 
 def paragraph_classification_batch_generator(

nucliadb/train/generators/paragraph_streaming.py

@@ -20,6 +20,8 @@
 
 from typing import AsyncGenerator
 
+from nidx_protos.nodereader_pb2 import StreamRequest
+
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.train import logger
@@ -29,7 +31,6 @@ from nucliadb_protos.dataset_pb2 import (
     ParagraphStreamItem,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamRequest
 
 
 def paragraph_streaming_batch_generator(

nucliadb/train/generators/question_answer_streaming.py

@@ -20,6 +20,8 @@
 
 from typing import AsyncGenerator
 
+from nidx_protos.nodereader_pb2 import StreamRequest
+
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR, FIELD_TYPE_STR_TO_PB
 from nucliadb.train import logger
@@ -33,7 +35,6 @@ from nucliadb_protos.dataset_pb2 import (
     QuestionAnswerStreamItem,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamRequest
 from nucliadb_protos.resources_pb2 import (
     FieldID,
     QuestionAnswer,

nucliadb/train/generators/sentence_classifier.py

@@ -21,6 +21,7 @@
 from typing import AsyncGenerator
 
 from fastapi import HTTPException
+from nidx_protos.nodereader_pb2 import StreamRequest
 
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
@@ -32,7 +33,6 @@ from nucliadb_protos.dataset_pb2 import (
     SentenceClassificationBatch,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamRequest
 
 
 def sentence_classification_batch_generator(

nucliadb/train/generators/token_classifier.py

@@ -21,6 +21,8 @@
 from collections import OrderedDict
 from typing import AsyncGenerator, cast
 
+from nidx_protos.nodereader_pb2 import StreamFilter, StreamRequest
+
 from nucliadb.common.cluster.base import AbstractIndexNode
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
 from nucliadb.train import logger
@@ -30,7 +32,6 @@ from nucliadb_protos.dataset_pb2 import (
     TokensClassification,
     TrainSet,
 )
-from nucliadb_protos.nodereader_pb2 import StreamFilter, StreamRequest
 
 NERS_DICT = dict[str, dict[str, list[tuple[int, int]]]]
 POSITION_DICT = OrderedDict[tuple[int, int], tuple[str, str]]

nucliadb/writer/back_pressure.py

@@ -25,7 +25,6 @@ from dataclasses import dataclass
 from datetime import datetime, timedelta
 from typing import Optional
 
-from async_lru import alru_cache
 from cachetools import TTLCache
 from fastapi import HTTPException, Request
 
@@ -442,29 +441,6 @@ def estimate_try_after(rate: float, pending: int, max_wait: int) -> datetime:
     return datetime.utcnow() + timedelta(seconds=delta_seconds)
 
 
-@alru_cache(maxsize=1024, ttl=60 * 15)
-async def get_nodes_for_kb_active_shards(context: ApplicationContext, kbid: str) -> list[str]:
-    with back_pressure_observer({"type": "get_kb_active_shard"}):
-        active_shard = await get_kb_active_shard(context, kbid)
-    if active_shard is None:
-        # KB doesn't exist or has been deleted
-        logger.debug("No active shard found for KB", extra={"kbid": kbid})
-        return []
-    return [replica.node for replica in active_shard.replicas]
-
-
-@alru_cache(maxsize=1024, ttl=60 * 60)
-async def get_nodes_for_resource_shard(
-    context: ApplicationContext, kbid: str, resource_uuid: str
-) -> list[str]:
-    with back_pressure_observer({"type": "get_resource_shard"}):
-        resource_shard = await get_resource_shard(context, kbid, resource_uuid)
-    if resource_shard is None:
-        # Resource doesn't exist or KB has been deleted
-        return []
-    return [replica.node for replica in resource_shard.replicas]
-
-
 async def get_nats_consumer_pending_messages(
     nats_manager: NatsConnectionManager, *, stream: str, consumer: str
 ) -> int:

{nucliadb-6.3.7.post4116.dist-info → nucliadb-6.3.7.post4120.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.3.7.post4116
+Version: 6.3.7.post4120
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.3.7.post4116
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.7.post4116
-Requires-Dist: nucliadb-protos>=6.3.7.post4116
-Requires-Dist: nucliadb-models>=6.3.7.post4116
-Requires-Dist: nidx-protos>=6.3.7.post4116
+Requires-Dist: nucliadb-telemetry[all]>=6.3.7.post4120
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.3.7.post4120
+Requires-Dist: nucliadb-protos>=6.3.7.post4120
+Requires-Dist: nucliadb-models>=6.3.7.post4120
+Requires-Dist: nidx-protos>=6.3.7.post4120
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]