nucliadb 6.3.7.post4081__py3-none-any.whl → 6.3.7.post4114__py3-none-any.whl

nucliadb/common/context/__init__.py

@@ -18,17 +18,19 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+from typing import Optional
 
 from nucliadb.common.cluster.manager import KBShardManager
 from nucliadb.common.cluster.settings import in_standalone_mode
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
-from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
+from nucliadb.common.nidx import NidxUtility, start_nidx_utility, stop_nidx_utility
 from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.partition import PartitionUtility
 from nucliadb_utils.settings import indexing_settings
 from nucliadb_utils.storages.storage import Storage
+from nucliadb_utils.transaction import TransactionUtility
 from nucliadb_utils.utilities import (
     get_storage,
     start_nats_manager,
@@ -42,16 +44,34 @@ from nucliadb_utils.utilities import (
 
 
 class ApplicationContext:
-    kv_driver: Driver
-    shard_manager: KBShardManager
-    blob_storage: Storage
-    partitioning: PartitionUtility
-    nats_manager: NatsConnectionManager
-
-    def __init__(self, service_name: str = "service") -> None:
+    def __init__(
+        self,
+        service_name: str = "service",
+        kv_driver: bool = True,
+        blob_storage: bool = True,
+        shard_manager: bool = True,
+        partitioning: bool = True,
+        nats_manager: bool = True,
+        transaction: bool = True,
+        nidx: bool = True,
+    ) -> None:
         self.service_name = service_name
         self._initialized: bool = False
         self._lock = asyncio.Lock()
+        self._kv_driver: Optional[Driver] = None
+        self._blob_storage: Optional[Storage] = None
+        self._shard_manager: Optional[KBShardManager] = None
+        self._partitioning: Optional[PartitionUtility] = None
+        self._nats_manager: Optional[NatsConnectionManager] = None
+        self._transaction: Optional[TransactionUtility] = None
+        self._nidx: Optional[NidxUtility] = None
+        self.enabled_kv_driver = kv_driver
+        self.enabled_blob_storage = blob_storage
+        self.enabled_shard_manager = shard_manager
+        self.enabled_partitioning = partitioning
+        self.enabled_nats_manager = nats_manager
+        self.enabled_transaction = transaction
+        self.enabled_nidx = nidx
 
     async def initialize(self) -> None:
         if self._initialized:
@@ -63,30 +83,75 @@ class ApplicationContext:
             self._initialized = True
 
     async def _initialize(self):
-        self.kv_driver = await setup_driver()
-        self.blob_storage = await get_storage()
-        self.shard_manager = await setup_cluster()
-        self.partitioning = start_partitioning_utility()
-        if not in_standalone_mode():
-            self.nats_manager = await start_nats_manager(
+        if self.enabled_kv_driver:
+            self._kv_driver = await setup_driver()
+        if self.enabled_blob_storage:
+            self._blob_storage = await get_storage()
+        if self.enabled_shard_manager:
+            self._shard_manager = await setup_cluster()
+        if self.enabled_partitioning:
+            self._partitioning = start_partitioning_utility()
+        if not in_standalone_mode() and self.enabled_nats_manager:
+            self._nats_manager = await start_nats_manager(
                 self.service_name,
                 indexing_settings.index_jetstream_servers,
                 indexing_settings.index_jetstream_auth,
             )
-        self.transaction = await start_transaction_utility(self.service_name)
-        self.nidx = await start_nidx_utility()
+        if self.enabled_transaction:
+            self._transaction = await start_transaction_utility(self.service_name)
+        if self.enabled_nidx:
+            self._nidx = await start_nidx_utility()
+
+    @property
+    def kv_driver(self) -> Driver:
+        assert self._kv_driver is not None, "Driver not initialized"
+        return self._kv_driver
+
+    @property
+    def shard_manager(self) -> KBShardManager:
+        assert self._shard_manager is not None, "Shard manager not initialized"
+        return self._shard_manager
+
+    @property
+    def blob_storage(self) -> Storage:
+        assert self._blob_storage is not None, "Blob storage not initialized"
+        return self._blob_storage
+
+    @property
+    def partitioning(self) -> PartitionUtility:
+        assert self._partitioning is not None, "Partitioning not initialized"
+        return self._partitioning
+
+    @property
+    def nats_manager(self) -> NatsConnectionManager:
+        assert self._nats_manager is not None, "NATS manager not initialized"
+        return self._nats_manager
+
+    @property
+    def transaction(self) -> TransactionUtility:
+        assert self._transaction is not None, "Transaction utility not initialized"
+        return self._transaction
+
+    @property
+    def nidx(self) -> NidxUtility:
+        assert self._nidx is not None, "Nidx utility not initialized"
+        return self._nidx
 
     async def finalize(self) -> None:
         if not self._initialized:
             return
-
-        await stop_nidx_utility()
-        await stop_transaction_utility()
-        if not in_standalone_mode():
+        if self.enabled_nidx:
+            await stop_nidx_utility()
+        if self.enabled_transaction:
+            await stop_transaction_utility()
+        if not in_standalone_mode() and self.enabled_nats_manager:
             await stop_nats_manager()
-
-        stop_partitioning_utility()
-        await teardown_cluster()
-        await teardown_driver()
-        await teardown_storage()
+        if self.enabled_partitioning:
+            stop_partitioning_utility()
+        if self.enabled_shard_manager:
+            await teardown_cluster()
+        if self.enabled_blob_storage:
+            await teardown_storage()
+        if self.enabled_kv_driver:
+            await teardown_driver()
         self._initialized = False

nucliadb/common/context/fastapi.py

@@ -19,6 +19,7 @@
 #
 
 from contextlib import asynccontextmanager
+from typing import Optional
 
 from fastapi import FastAPI
 from starlette.routing import Mount
@@ -27,8 +28,9 @@ from nucliadb.common.context import ApplicationContext
 
 
 @asynccontextmanager
-async def inject_app_context(app: FastAPI):
-    context = ApplicationContext()
+async def inject_app_context(app: FastAPI, context: Optional[ApplicationContext] = None):
+    if context is None:
+        context = ApplicationContext()
 
     app.state.context = context
 

nucliadb/ingest/consumer/consumer.py

@@ -160,6 +160,8 @@ class IngestConsumer:
                 logger.warning("Could not delete blob reference", exc_info=True)
 
     async def subscription_worker(self, msg: Msg):
+        context.clear_context()
+
         kbid: Optional[str] = None
         subject = msg.subject
         reply = msg.reply
@@ -182,7 +184,6 @@ class IngestConsumer:
             MessageProgressUpdater(msg, nats_consumer_settings.nats_ack_wait * 0.66),
             self.lock,
         ):
-            logger.info(f"Message processing: subject:{subject}, seqid: {seqid}, reply: {reply}")
             try:
                 pb = await self.get_broker_message(msg)
                 if pb.source == pb.MessageSource.PROCESSOR:
@@ -194,10 +195,8 @@ class IngestConsumer:
                 else:
                     audit_time = ""
 
-                logger.debug(
-                    f"Received from {message_source} on {pb.kbid}/{pb.uuid} seq {seqid} partition {self.partition} at {time}"  # noqa
-                )
                 context.add_context({"kbid": pb.kbid, "rid": pb.uuid})
+                logger.info(f"Message processing: subject:{subject}, seqid: {seqid}, reply: {reply}")
                 kbid = pb.kbid
                 try:
                     source = "writer" if pb.source == pb.MessageSource.WRITER else "processor"

nucliadb/search/api/v1/find.py

@@ -40,6 +40,7 @@ from nucliadb_models.configuration import FindConfig
 from nucliadb_models.filters import FilterExpression
 from nucliadb_models.resource import ExtractedDataTypeName, NucliaDBRoles
 from nucliadb_models.search import (
+    FindOptions,
     FindRequest,
     KnowledgeboxFindResults,
     NucliaDBClientType,
@@ -47,7 +48,6 @@ from nucliadb_models.search import (
     Reranker,
     RerankerName,
     ResourceProperties,
-    SearchOptions,
     SearchParamDefaults,
 )
 from nucliadb_models.security import RequestSecurity
@@ -61,7 +61,7 @@ FIND_EXAMPLES = {
         description="Perform a hybrid search that will return text and semantic results matching the query",
         value={
             "query": "How can I be an effective product manager?",
-            "features": [SearchOptions.KEYWORD, SearchOptions.SEMANTIC],
+            "features": [FindOptions.KEYWORD, FindOptions.SEMANTIC],
         },
     )
 }
@@ -110,11 +110,11 @@ async def find_knowledgebox(
     range_modification_end: Optional[DateTime] = fastapi_query(
         SearchParamDefaults.range_modification_end
     ),
-    features: list[SearchOptions] = fastapi_query(
+    features: list[FindOptions] = fastapi_query(
         SearchParamDefaults.search_features,
         default=[
-            SearchOptions.KEYWORD,
-            SearchOptions.SEMANTIC,
+            FindOptions.KEYWORD,
+            FindOptions.SEMANTIC,
         ],
     ),
     debug: bool = fastapi_query(SearchParamDefaults.debug),

nucliadb/search/api/v1/search.py

@@ -37,11 +37,9 @@ from nucliadb.search.search import cache
 from nucliadb.search.search.exceptions import InvalidQueryError
 from nucliadb.search.search.merge import merge_results
 from nucliadb.search.search.query_parser.parsers.search import parse_search
-from nucliadb.search.search.query_parser.parsers.unit_retrieval import convert_retrieval_to_proto
+from nucliadb.search.search.query_parser.parsers.unit_retrieval import legacy_convert_retrieval_to_proto
 from nucliadb.search.search.utils import (
-    min_score_from_payload,
     min_score_from_query_params,
-    should_disable_vector_search,
 )
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.filters import FilterExpression
@@ -263,14 +261,8 @@ async def search(
     audit = get_audit()
     start_time = time()
 
-    item.min_score = min_score_from_payload(item.min_score)
-
-    if SearchOptions.SEMANTIC in item.features:
-        if should_disable_vector_search(item):
-            item.features.remove(SearchOptions.SEMANTIC)
-
     parsed = await parse_search(kbid, item)
-    pb_query, incomplete_results, autofilters, _ = await convert_retrieval_to_proto(parsed)
+    pb_query, incomplete_results, autofilters, _ = await legacy_convert_retrieval_to_proto(parsed)
 
     # We need to query all nodes
     results, query_incomplete_results, queried_nodes = await node_query(kbid, Method.SEARCH, pb_query)

nucliadb/search/search/chat/ask.py

@@ -80,6 +80,7 @@ from nucliadb_models.search import (
     CitationsAskResponseItem,
     DebugAskResponseItem,
     ErrorAskResponseItem,
+    FindOptions,
     FindParagraph,
     FindRequest,
     GraphStrategy,
@@ -97,7 +98,6 @@ from nucliadb_models.search import (
     Relations,
     RelationsAskResponseItem,
     RetrievalAskResponseItem,
-    SearchOptions,
     StatusAskResponseItem,
     SyncAskMetadata,
     SyncAskResponse,
@@ -755,6 +755,9 @@ async def retrieval_in_kb(
         )
 
         if graph_strategy is not None:
+            assert parsed_query.retrieval.reranker is not None, (
+                "find parser must provide a reranking algorithm"
+            )
             reranker = get_reranker(parsed_query.retrieval.reranker)
             graph_results, graph_request = await get_graph_results(
                 kbid=kbid,
@@ -952,9 +955,9 @@ def calculate_prequeries_for_json_schema(
     json_schema = ask_request.answer_json_schema or {}
     features = []
     if ChatOptions.SEMANTIC in ask_request.features:
-        features.append(SearchOptions.SEMANTIC)
+        features.append(FindOptions.SEMANTIC)
     if ChatOptions.KEYWORD in ask_request.features:
-        features.append(SearchOptions.KEYWORD)
+        features.append(FindOptions.KEYWORD)
 
     properties = json_schema.get("parameters", {}).get("properties", {})
     if len(properties) == 0:  # pragma: no cover

nucliadb/search/search/chat/query.py

@@ -29,7 +29,8 @@ from nucliadb.search.search.exceptions import IncompleteFindResultsError
 from nucliadb.search.search.find import find
 from nucliadb.search.search.merge import merge_relations_results
 from nucliadb.search.search.metrics import RAGMetrics
-from nucliadb.search.search.query_parser.models import ParsedQuery
+from nucliadb.search.search.query_parser.models import ParsedQuery, Query, RelationQuery, UnitRetrieval
+from nucliadb.search.search.query_parser.parsers.unit_retrieval import convert_retrieval_to_proto
 from nucliadb.search.settings import settings
 from nucliadb.search.utilities import get_predict
 from nucliadb_models import filters
@@ -37,6 +38,7 @@ from nucliadb_models.search import (
     AskRequest,
     ChatContextMessage,
     ChatOptions,
+    FindOptions,
     FindRequest,
     KnowledgeboxFindResults,
     NucliaDBClientType,
@@ -47,14 +49,11 @@ from nucliadb_models.search import (
     PromptContextOrder,
     Relations,
     RephraseModel,
-    SearchOptions,
     parse_rephrase_prompt,
 )
 from nucliadb_protos import audit_pb2
 from nucliadb_protos.nodereader_pb2 import (
-    EntitiesSubgraphRequest,
-    RelationSearchResponse,
-    SearchRequest,
+    GraphSearchResponse,
     SearchResponse,
 )
 from nucliadb_protos.utils_pb2 import RelationNode
@@ -181,11 +180,11 @@ def find_request_from_ask_request(item: AskRequest, query: str) -> FindRequest:
     find_request.resource_filters = item.resource_filters
     find_request.features = []
     if ChatOptions.SEMANTIC in item.features:
-        find_request.features.append(SearchOptions.SEMANTIC)
+        find_request.features.append(FindOptions.SEMANTIC)
     if ChatOptions.KEYWORD in item.features:
-        find_request.features.append(SearchOptions.KEYWORD)
+        find_request.features.append(FindOptions.KEYWORD)
     if ChatOptions.RELATIONS in item.features:
-        find_request.features.append(SearchOptions.RELATIONS)
+        find_request.features.append(FindOptions.RELATIONS)
     find_request.query = query
     find_request.fields = item.fields
     find_request.filters = item.filters
@@ -274,13 +273,18 @@ async def get_relations_results_from_entities(
     only_entity_to_entity: bool = False,
     deleted_entities: set[str] = set(),
 ) -> Relations:
-    request = SearchRequest()
-    request.relation_subgraph.entry_points.extend(entities)
-    request.relation_subgraph.depth = 1
-
-    deleted = EntitiesSubgraphRequest.DeletedEntities()
-    deleted.node_values.extend(deleted_entities)
-    request.relation_subgraph.deleted_entities.append(deleted)
+    entry_points = list(entities)
+    retrieval = UnitRetrieval(
+        query=Query(
+            relation=RelationQuery(
+                entry_points=entry_points,
+                deleted_entities={"": list(deleted_entities)},
+                deleted_entity_groups=[],
+            )
+        ),
+        top_k=50,
+    )
+    request = convert_retrieval_to_proto(retrieval)
 
     results: list[SearchResponse]
     (
@@ -293,10 +297,10 @@ async def get_relations_results_from_entities(
         request,
         timeout=timeout,
     )
-    relations_results: list[RelationSearchResponse] = [result.relation for result in results]
+    relations_results: list[GraphSearchResponse] = [result.graph for result in results]
     return await merge_relations_results(
         relations_results,
-        request.relation_subgraph.entry_points,
+        entry_points,
         only_with_metadata,
         only_agentic_relations,
         only_entity_to_entity,

nucliadb/search/search/find.py

@@ -38,7 +38,7 @@ from nucliadb.search.search.metrics import (
 )
 from nucliadb.search.search.query_parser.models import ParsedQuery
 from nucliadb.search.search.query_parser.parsers import parse_find
-from nucliadb.search.search.query_parser.parsers.unit_retrieval import convert_retrieval_to_proto
+from nucliadb.search.search.query_parser.parsers.unit_retrieval import legacy_convert_retrieval_to_proto
 from nucliadb.search.search.rank_fusion import (
     get_rank_fusion,
 )
@@ -92,11 +92,17 @@ async def _index_node_retrieval(
 
     with metrics.time("query_parse"):
         parsed = await parse_find(kbid, item)
+        assert parsed.retrieval.rank_fusion is not None and parsed.retrieval.reranker is not None, (
+            "find parser must provide rank fusion and reranker algorithms"
+        )
         rank_fusion = get_rank_fusion(parsed.retrieval.rank_fusion)
         reranker = get_reranker(parsed.retrieval.reranker)
-        pb_query, incomplete_results, autofilters, rephrased_query = await convert_retrieval_to_proto(
-            parsed
-        )
+        (
+            pb_query,
+            incomplete_results,
+            autofilters,
+            rephrased_query,
+        ) = await legacy_convert_retrieval_to_proto(parsed)
 
     with metrics.time("node_query"):
         results, query_incomplete_results, queried_nodes = await node_query(
@@ -181,8 +187,11 @@ async def _external_index_retrieval(
     """
     # Parse query
     parsed = await parse_find(kbid, item)
+    assert parsed.retrieval.reranker is not None, "find parser must provide a reranking algorithm"
     reranker = get_reranker(parsed.retrieval.reranker)
-    search_request, incomplete_results, _, rephrased_query = await convert_retrieval_to_proto(parsed)
+    search_request, incomplete_results, _, rephrased_query = await legacy_convert_retrieval_to_proto(
+        parsed
+    )
 
     # Query index
     query_results = await external_index_manager.query(search_request)  # noqa

nucliadb/search/search/find_merge.py

@@ -52,9 +52,9 @@ from nucliadb_models.search import (
 )
 from nucliadb_protos.nodereader_pb2 import (
     DocumentScored,
+    GraphSearchResponse,
     ParagraphResult,
     ParagraphSearchResponse,
-    RelationSearchResponse,
     SearchResponse,
     VectorSearchResponse,
 )
@@ -142,8 +142,8 @@ async def build_find_response(
     # build relations graph
     entry_points = []
     if retrieval.query.relation is not None:
-        entry_points = retrieval.query.relation.detected_entities
-    relations = await merge_relations_results([search_response.relation], entry_points)
+        entry_points = retrieval.query.relation.entry_points
+    relations = await merge_relations_results([search_response.graph], entry_points)
 
     # compose response
     find_resources = compose_find_resources(text_blocks, resources)
@@ -178,16 +178,16 @@ def merge_shard_responses(
     """
     paragraphs = []
     vectors = []
-    relations = []
+    graphs = []
     for response in responses:
         paragraphs.append(response.paragraph)
         vectors.append(response.vector)
-        relations.append(response.relation)
+        graphs.append(response.graph)
 
     merged = SearchResponse(
         paragraph=merge_shards_keyword_responses(paragraphs),
         vector=merge_shards_semantic_responses(vectors),
-        relation=merge_shards_relation_responses(relations),
+        graph=merge_shards_graph_responses(graphs),
     )
     return merged
 
@@ -230,13 +230,27 @@ def merge_shards_semantic_responses(
     return merged
 
 
-def merge_shards_relation_responses(
-    relation_responses: list[RelationSearchResponse],
-) -> RelationSearchResponse:
-    merged = RelationSearchResponse()
-    for response in relation_responses:
-        merged.prefix.nodes.extend(response.prefix.nodes)
-        merged.subgraph.relations.extend(response.subgraph.relations)
+def merge_shards_graph_responses(
+    graph_responses: list[GraphSearchResponse],
+):
+    merged = GraphSearchResponse()
+
+    for response in graph_responses:
+        nodes_offset = len(merged.nodes)
+        relations_offset = len(merged.relations)
+
+        # paths contain indexes to nodes and relations, we must offset them
+        # while merging responses to maintain valid data
+        for path in response.graph:
+            merged_path = GraphSearchResponse.Path()
+            merged_path.CopyFrom(path)
+            merged_path.source += nodes_offset
+            merged_path.relation += relations_offset
+            merged_path.destination += nodes_offset
+            merged.graph.append(merged_path)
+
+        merged.nodes.extend(response.nodes)
+        merged.relations.extend(response.relations)
 
     return merged
 

nucliadb/search/search/merge.py

@@ -65,9 +65,9 @@ from nucliadb_protos.nodereader_pb2 import (
     DocumentResult,
     DocumentScored,
     DocumentSearchResponse,
+    GraphSearchResponse,
     ParagraphResult,
     ParagraphSearchResponse,
-    RelationSearchResponse,
     SearchResponse,
     SuggestResponse,
     VectorSearchResponse,
@@ -438,7 +438,7 @@ async def merge_paragraph_results(
 
 @merge_observer.wrap({"type": "merge_relations"})
 async def merge_relations_results(
-    relations_responses: list[RelationSearchResponse],
+    graph_responses: list[GraphSearchResponse],
     query_entry_points: Iterable[RelationNode],
     only_with_metadata: bool = False,
     only_agentic: bool = False,
@@ -448,7 +448,7 @@ async def merge_relations_results(
     return await loop.run_in_executor(
         None,
         _merge_relations_results,
-        relations_responses,
+        graph_responses,
         query_entry_points,
         only_with_metadata,
         only_agentic,
@@ -457,7 +457,7 @@ async def merge_relations_results(
 
 
 def _merge_relations_results(
-    relations_responses: list[RelationSearchResponse],
+    graph_responses: list[GraphSearchResponse],
     query_entry_points: Iterable[RelationNode],
     only_with_metadata: bool,
     only_agentic: bool,
@@ -480,17 +480,16 @@ def _merge_relations_results(
     for entry_point in query_entry_points:
         relations.entities[entry_point.value] = EntitySubgraph(related_to=[])
 
-    for relation_response in relations_responses:
-        for index_relation in relation_response.subgraph.relations:
-            relation = index_relation.relation
-            origin = relation.source
-            destination = relation.to
-            relation_type = RelationTypePbMap[relation.relation]  # type: ignore
-            relation_label = relation.relation_label
-            metadata = relation.metadata if relation.HasField("metadata") else None
-
-            if index_relation.resource_field_id is not None:
-                resource_id = index_relation.resource_field_id.split("/")[0]
+    for graph_response in graph_responses:
+        for path in graph_response.graph:
+            relation = graph_response.relations[path.relation]
+            origin = graph_response.nodes[path.source]
+            destination = graph_response.nodes[path.destination]
+            relation_type = RelationTypePbMap[relation.relation_type]
+            relation_label = relation.label
+            metadata = path.metadata if path.HasField("metadata") else None
+            if path.resource_field_id is not None:
+                resource_id = path.resource_field_id.split("/")[0]
 
             # If only_with_metadata is True, we check that metadata for the relation is not None
             # If only_agentic is True, we check that metadata for the relation is not None and that it has a data_augmentation_task_id
@@ -547,13 +546,13 @@ async def merge_results(
     paragraphs = []
     documents = []
     vectors = []
-    relations = []
+    graphs = []
 
     for response in search_responses:
         paragraphs.append(response.paragraph)
         documents.append(response.document)
         vectors.append(response.vector)
-        relations.append(response.relation)
+        graphs.append(response.graph)
 
     api_results = KnowledgeboxSearchResults()
 
@@ -595,7 +594,7 @@ async def merge_results(
 
     if retrieval.query.relation is not None:
         api_results.relations = await merge_relations_results(
-            relations, retrieval.query.relation.detected_entities
+            graphs, retrieval.query.relation.entry_points
         )
 
     api_results.resources = await fetch_resources(resources, kbid, show, field_type_filter, extracted)