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

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)

nucliadb/search/search/query_parser/models.py

@@ -21,10 +21,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Literal, Optional, Union
 
-from pydantic import (
-    BaseModel,
-    Field,
-)
+from pydantic import BaseModel, ConfigDict, Field
 
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb_models import search as search_models
@@ -35,8 +32,7 @@ from nucliadb_protos import nodereader_pb2, utils_pb2
 # query
 
 
-@dataclass
-class _TextQuery:
+class _TextQuery(BaseModel):
     query: str
     is_synonyms_query: bool
     min_score: float
@@ -48,24 +44,23 @@ FulltextQuery = _TextQuery
 KeywordQuery = _TextQuery
 
 
-@dataclass
-class SemanticQuery:
+class SemanticQuery(BaseModel):
     query: Optional[list[float]]
     vectorset: str
     min_score: float
 
 
-@dataclass
-class RelationQuery:
-    detected_entities: list[utils_pb2.RelationNode]
+class RelationQuery(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    entry_points: list[utils_pb2.RelationNode]
     # list[subtype]
     deleted_entity_groups: list[str]
     # subtype -> list[entity]
     deleted_entities: dict[str, list[str]]
 
 
-@dataclass
-class Query:
+class Query(BaseModel):
     fulltext: Optional[FulltextQuery] = None
     keyword: Optional[KeywordQuery] = None
     semantic: Optional[SemanticQuery] = None
@@ -75,8 +70,9 @@ class Query:
 # filters
 
 
-@dataclass
-class Filters:
+class Filters(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     field_expression: Optional[nodereader_pb2.FilterExpression] = None
     paragraph_expression: Optional[nodereader_pb2.FilterExpression] = None
     filter_expression_operator: nodereader_pb2.FilterOperator.ValueType = (
@@ -125,30 +121,29 @@ Reranker = Union[NoopReranker, PredictReranker]
 # retrieval and generation operations
 
 
-@dataclass
-class UnitRetrieval:
+class UnitRetrieval(BaseModel):
     query: Query
     top_k: int
-    filters: Filters
-    # TODO: rank fusion depends on the response building, not the retrieval
-    rank_fusion: RankFusion
-    # TODO: reranking fusion depends on the response building, not the retrieval
-    reranker: Reranker
+    filters: Filters = Field(default_factory=Filters)
+    rank_fusion: Optional[RankFusion] = None
+    reranker: Optional[Reranker] = None
 
 
-@dataclass
-class Generation:
+# TODO: augmentation things: hydration...
+
+
+class Generation(BaseModel):
     use_visual_llm: bool
     max_context_tokens: int
     max_answer_tokens: Optional[int]
 
 
-@dataclass
-class ParsedQuery:
+class ParsedQuery(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     fetcher: Fetcher
     retrieval: UnitRetrieval
     generation: Optional[Generation] = None
-    # TODO: add merge, rank fusion, rerank...
 
 
 ### Catalog

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

@@ -28,7 +28,6 @@ from nucliadb.search.search.query_parser.models import (
     KeywordQuery,
     SemanticQuery,
 )
-from nucliadb.search.search.utils import should_disable_vector_search
 from nucliadb_models import search as search_models
 
 DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
@@ -38,28 +37,40 @@ DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
 INVALID_QUERY = re.compile(r"- +\*")
 
 
-def validate_base_request(item: search_models.BaseSearchRequest):
+def validate_query_syntax(query: str):
     # Filter some queries that panic tantivy, better than returning the 500
-    if INVALID_QUERY.search(item.query):
+    if INVALID_QUERY.search(query):
         raise InvalidQueryError("query", "Invalid query syntax")
 
-    # synonyms are not compatible with vector/graph search
-    if (
-        item.with_synonyms
-        and item.query
-        and (
-            search_models.SearchOptions.SEMANTIC in item.features
-            or search_models.SearchOptions.RELATIONS in item.features
-        )
-    ):
-        raise InvalidQueryError(
-            "synonyms",
-            "Search with custom synonyms is only supported on paragraph and document search",
-        )
-
-    if search_models.SearchOptions.SEMANTIC in item.features:
-        if should_disable_vector_search(item):
-            item.features.remove(search_models.SearchOptions.SEMANTIC)
+
+def is_empty_query(request: search_models.BaseSearchRequest) -> bool:
+    return len(request.query) == 0
+
+
+def has_user_vectors(request: search_models.BaseSearchRequest) -> bool:
+    return request.vector is not None and len(request.vector) > 0
+
+
+def is_exact_match_only_query(request: search_models.BaseSearchRequest) -> bool:
+    """
+    '"something"' -> True
+    'foo "something" else' -> False
+    """
+    query = request.query.strip()
+    return len(query) > 0 and query.startswith('"') and query.endswith('"')
+
+
+def should_disable_vector_search(request: search_models.BaseSearchRequest) -> bool:
+    if has_user_vectors(request):
+        return False
+
+    if is_exact_match_only_query(request):
+        return True
+
+    if is_empty_query(request):
+        return True
+
+    return False
 
 
 def parse_top_k(item: search_models.BaseSearchRequest) -> int:
@@ -92,7 +103,7 @@ async def parse_keyword_query(
 
 
 async def parse_semantic_query(
-    item: search_models.BaseSearchRequest,
+    item: Union[search_models.SearchRequest, search_models.FindRequest],
     *,
     fetcher: Fetcher,
 ) -> SemanticQuery: