nucliadb 6.3.6.post4063__py3-none-any.whl → 6.3.7.post4068__py3-none-any.whl

nucliadb/search/api/v1/search.py

@@ -36,10 +36,9 @@ from nucliadb.search.requesters.utils import Method, debug_nodes_info, node_quer
 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 import QueryParser
-from nucliadb.search.search.query_parser.old_filters import OldFilterParams
+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.utils import (
-    filter_hidden_resources,
     min_score_from_payload,
     min_score_from_query_params,
     should_disable_vector_search,
@@ -270,53 +269,21 @@ async def search(
         if should_disable_vector_search(item):
             item.features.remove(SearchOptions.SEMANTIC)
 
-    # We need to query all nodes
-    query_parser = QueryParser(
-        kbid=kbid,
-        features=item.features,
-        query=item.query,
-        filter_expression=item.filter_expression,
-        faceted=item.faceted,
-        sort=item.sort,
-        top_k=item.top_k,
-        min_score=item.min_score,
-        old_filters=OldFilterParams(
-            label_filters=item.filters,
-            keyword_filters=[],
-            range_creation_start=item.range_creation_start,
-            range_creation_end=item.range_creation_end,
-            range_modification_start=item.range_modification_start,
-            range_modification_end=item.range_modification_end,
-            fields=item.fields,
-        ),
-        user_vector=item.vector,
-        vectorset=item.vectorset,
-        with_duplicates=item.with_duplicates,
-        with_status=with_status,
-        with_synonyms=item.with_synonyms,
-        autofilter=item.autofilter,
-        security=item.security,
-        rephrase=item.rephrase,
-        hidden=await filter_hidden_resources(kbid, item.show_hidden),
-        rephrase_prompt=item.rephrase_prompt,
-    )
-    pb_query, incomplete_results, autofilters, _ = await query_parser.parse()
+    parsed = await parse_search(kbid, item)
+    pb_query, incomplete_results, autofilters, _ = await 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)
-
     incomplete_results = incomplete_results or query_incomplete_results
 
     # We need to merge
     search_results = await merge_results(
         results,
-        top_k=item.top_k,
+        parsed.retrieval,
         kbid=kbid,
         show=item.show,
         field_type_filter=item.field_type_filter,
         extracted=item.extracted,
-        sort=query_parser.sort,  # type: ignore
-        requested_relations=pb_query.relation_subgraph,
-        min_score=query_parser.min_score,
         highlight=item.highlight,
     )
 

nucliadb/search/search/chat/ask.py

@@ -61,8 +61,11 @@ from nucliadb.search.search.exceptions import (
 )
 from nucliadb.search.search.graph_strategy import get_graph_results
 from nucliadb.search.search.metrics import RAGMetrics
-from nucliadb.search.search.query import QueryParser
-from nucliadb.search.search.query_parser.old_filters import OldFilterParams
+from nucliadb.search.search.query_parser.fetcher import Fetcher
+from nucliadb.search.search.query_parser.parsers.ask import fetcher_for_ask, parse_ask
+from nucliadb.search.search.rerankers import (
+    get_reranker,
+)
 from nucliadb.search.utilities import get_predict
 from nucliadb_models.search import (
     AnswerAskResponseItem,
@@ -83,7 +86,6 @@ from nucliadb_models.search import (
     JSONAskResponseItem,
     KnowledgeboxFindResults,
     MetadataAskResponseItem,
-    MinScore,
     NucliaDBClientType,
     PrequeriesAskResponseItem,
     PreQueriesStrategy,
@@ -116,7 +118,7 @@ class RetrievalMatch:
 @dataclasses.dataclass
 class RetrievalResults:
     main_query: KnowledgeboxFindResults
-    query_parser: QueryParser
+    fetcher: Fetcher
     main_query_weight: float
     prequeries: Optional[list[PreQueryResult]] = None
     best_matches: list[RetrievalMatch] = dataclasses.field(default_factory=list)
@@ -543,12 +545,12 @@ async def ask(
             prequeries_results=err.prequeries,
         )
 
-    query_parser = retrieval_results.query_parser
+    # parse ask request generation parameters reusing the same fetcher as
+    # retrieval, to avoid multiple round trips to Predict API
+    generation = await parse_ask(kbid, ask_request, fetcher=retrieval_results.fetcher)
 
     # Now we build the prompt context
     with metrics.time("context_building"):
-        query_parser.max_tokens = ask_request.max_tokens  # type: ignore
-        max_tokens_context = await query_parser.get_max_tokens_context()
         prompt_context_builder = PromptContextBuilder(
             kbid=kbid,
             ordered_paragraphs=[match.paragraph for match in retrieval_results.best_matches],
@@ -557,8 +559,8 @@ async def ask(
             user_image_context=ask_request.extra_context_images,
             strategies=ask_request.rag_strategies,
             image_strategies=ask_request.rag_images_strategies,
-            max_context_characters=tokens_to_chars(max_tokens_context),
-            visual_llm=await query_parser.get_visual_llm_enabled(),
+            max_context_characters=tokens_to_chars(generation.max_context_tokens),
+            visual_llm=generation.use_visual_llm,
         )
         (
             prompt_context,
@@ -580,7 +582,7 @@ async def ask(
         citations=ask_request.citations,
         citation_threshold=ask_request.citation_threshold,
         generative_model=ask_request.generative_model,
-        max_tokens=query_parser.get_max_tokens_answer(),
+        max_tokens=generation.max_answer_tokens,
         query_context_images=prompt_context_images,
         json_schema=ask_request.answer_json_schema,
         rerank_context=False,
@@ -741,7 +743,7 @@ async def retrieval_in_kb(
     prequeries = parse_prequeries(ask_request)
     graph_strategy = parse_graph_strategy(ask_request)
     with metrics.time("retrieval"):
-        main_results, prequeries_results, query_parser = await get_find_results(
+        main_results, prequeries_results, parsed_query = await get_find_results(
             kbid=kbid,
             query=main_query,
             item=ask_request,
@@ -753,6 +755,7 @@ async def retrieval_in_kb(
         )
 
         if graph_strategy is not None:
+            reranker = get_reranker(parsed_query.retrieval.reranker)
             graph_results, graph_request = await get_graph_results(
                 kbid=kbid,
                 query=main_query,
@@ -762,6 +765,7 @@ async def retrieval_in_kb(
                 origin=origin,
                 graph_strategy=graph_strategy,
                 metrics=metrics,
+                text_block_reranker=reranker,
             )
 
             if prequeries_results is None:
@@ -784,7 +788,7 @@ async def retrieval_in_kb(
     return RetrievalResults(
         main_query=main_results,
         prequeries=prequeries_results,
-        query_parser=query_parser,
+        fetcher=parsed_query.fetcher,
         main_query_weight=main_query_weight,
         best_matches=best_matches,
     )
@@ -805,18 +809,7 @@ async def retrieval_in_resource(
         return RetrievalResults(
             main_query=KnowledgeboxFindResults(resources={}, min_score=None),
             prequeries=None,
-            query_parser=QueryParser(
-                kbid=kbid,
-                features=[],
-                query="",
-                filter_expression=ask_request.filter_expression,
-                old_filters=OldFilterParams(
-                    label_filters=ask_request.filters,
-                    keyword_filters=ask_request.keyword_filters,
-                ),
-                top_k=0,
-                min_score=MinScore(),
-            ),
+            fetcher=fetcher_for_ask(kbid, ask_request),
             main_query_weight=1.0,
         )
 
@@ -836,7 +829,7 @@ async def retrieval_in_resource(
             add_resource_filter(prequery.request, [resource])
 
     with metrics.time("retrieval"):
-        main_results, prequeries_results, query_parser = await get_find_results(
+        main_results, prequeries_results, parsed_query = await get_find_results(
             kbid=kbid,
             query=main_query,
             item=ask_request,
@@ -859,7 +852,7 @@ async def retrieval_in_resource(
     return RetrievalResults(
         main_query=main_results,
         prequeries=prequeries_results,
-        query_parser=query_parser,
+        fetcher=parsed_query.fetcher,
         main_query_weight=main_query_weight,
         best_matches=best_matches,
     )

nucliadb/search/search/chat/query.py

@@ -29,7 +29,7 @@ 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 import QueryParser
+from nucliadb.search.search.query_parser.models import ParsedQuery
 from nucliadb.search.settings import settings
 from nucliadb.search.utilities import get_predict
 from nucliadb_models import filters
@@ -93,7 +93,7 @@ async def get_find_results(
     origin: str,
     metrics: RAGMetrics = RAGMetrics(),
     prequeries_strategy: Optional[PreQueriesStrategy] = None,
-) -> tuple[KnowledgeboxFindResults, Optional[list[PreQueryResult]], QueryParser]:
+) -> tuple[KnowledgeboxFindResults, Optional[list[PreQueryResult]], ParsedQuery]:
     prequeries_results = None
     prefilter_queries_results = None
     queries_results = None
@@ -223,10 +223,10 @@ async def run_main_query(
     user: str,
     origin: str,
     metrics: RAGMetrics = RAGMetrics(),
-) -> tuple[KnowledgeboxFindResults, QueryParser]:
+) -> tuple[KnowledgeboxFindResults, ParsedQuery]:
     find_request = find_request_from_ask_request(item, query)
 
-    find_results, incomplete, query_parser = await find(
+    find_results, incomplete, parsed_query = await find(
         kbid,
         find_request,
         ndb_client,
@@ -237,7 +237,7 @@ async def run_main_query(
     )
     if incomplete:
         raise IncompleteFindResultsError()
-    return find_results, query_parser
+    return find_results, parsed_query
 
 
 async def get_relations_results(
@@ -297,7 +297,7 @@ async def get_relations_results_from_entities(
     relations_results: list[RelationSearchResponse] = [result.relation for result in results]
     return await merge_relations_results(
         relations_results,
-        request.relation_subgraph,
+        request.relation_subgraph.entry_points,
         only_with_metadata,
         only_agentic_relations,
         only_entity_to_entity,

nucliadb/search/search/find.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from dataclasses import dataclass
 from time import time
 from typing import Optional
 
@@ -38,30 +37,22 @@ from nucliadb.search.search.hydrator import (
 from nucliadb.search.search.metrics import (
     RAGMetrics,
 )
-from nucliadb.search.search.query import QueryParser
-from nucliadb.search.search.query_parser.old_filters import OldFilterParams
+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.rank_fusion import (
-    RankFusionAlgorithm,
     get_rank_fusion,
 )
 from nucliadb.search.search.rerankers import (
-    Reranker,
     RerankingOptions,
     get_reranker,
 )
-from nucliadb.search.search.utils import (
-    filter_hidden_resources,
-    min_score_from_payload,
-    should_disable_vector_search,
-)
 from nucliadb.search.settings import settings
 from nucliadb_models.search import (
     FindRequest,
     KnowledgeboxFindResults,
     MinScore,
     NucliaDBClientType,
-    SearchOptions,
 )
 from nucliadb_utils.utilities import get_audit
 
@@ -76,7 +67,7 @@ async def find(
     x_forwarded_for: str,
     generative_model: Optional[str] = None,
     metrics: RAGMetrics = RAGMetrics(),
-) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
+) -> tuple[KnowledgeboxFindResults, bool, ParsedQuery]:
     external_index_manager = await get_external_index_manager(kbid=kbid)
     if external_index_manager is not None:
         return await _external_index_retrieval(
@@ -99,15 +90,17 @@ async def _index_node_retrieval(
     x_forwarded_for: str,
     generative_model: Optional[str] = None,
     metrics: RAGMetrics = RAGMetrics(),
-) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
+) -> tuple[KnowledgeboxFindResults, bool, ParsedQuery]:
     audit = get_audit()
     start_time = time()
 
-    query_parser, rank_fusion, reranker = await query_parser_from_find_request(
-        kbid, item, generative_model=generative_model
-    )
     with metrics.time("query_parse"):
-        pb_query, incomplete_results, autofilters, rephrased_query = await query_parser.parse()
+        parsed = await parse_find(kbid, item, generative_model=generative_model)
+        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
+        )
 
     with metrics.time("node_query"):
         results, query_incomplete_results, queried_nodes = await node_query(
@@ -119,13 +112,10 @@ async def _index_node_retrieval(
     with metrics.time("results_merge"):
         search_results = await build_find_response(
             results,
+            retrieval=parsed.retrieval,
             kbid=kbid,
             query=pb_query.body,
             rephrased_query=rephrased_query,
-            relation_subgraph_query=pb_query.relation_subgraph,
-            min_score_bm25=pb_query.min_score_bm25,
-            min_score_semantic=pb_query.min_score_semantic,
-            top_k=item.top_k,
             show=item.show,
             extracted=item.extracted,
             field_type_filter=item.field_type_filter,
@@ -182,7 +172,7 @@ async def _index_node_retrieval(
             },
         )
 
-    return search_results, incomplete_results, query_parser
+    return search_results, incomplete_results, parsed
 
 
 async def _external_index_retrieval(
@@ -190,15 +180,14 @@ async def _external_index_retrieval(
     item: FindRequest,
     external_index_manager: ExternalIndexManager,
     generative_model: Optional[str] = None,
-) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
+) -> tuple[KnowledgeboxFindResults, bool, ParsedQuery]:
     """
     Parse the query, query the external index, and hydrate the results.
     """
     # Parse query
-    query_parser, _, reranker = await query_parser_from_find_request(
-        kbid, item, generative_model=generative_model
-    )
-    search_request, incomplete_results, _, rephrased_query = await query_parser.parse()
+    parsed = await parse_find(kbid, item, generative_model=generative_model)
+    reranker = get_reranker(parsed.retrieval.reranker)
+    search_request, incomplete_results, _, rephrased_query = await convert_retrieval_to_proto(parsed)
 
     # Query index
     query_results = await external_index_manager.query(search_request)  # noqa
@@ -218,13 +207,15 @@ async def _external_index_retrieval(
             kbid=kbid,
             query=search_request.body,
         ),
-        top_k=query_parser.top_k,
+        top_k=parsed.retrieval.top_k,
     )
     find_resources = compose_find_resources(text_blocks, resources)
 
     results_min_score = MinScore(
         bm25=0,
-        semantic=query_parser.min_score.semantic,
+        semantic=parsed.retrieval.query.semantic.min_score
+        if parsed.retrieval.query.semantic is not None
+        else 0.0,
     )
     retrieval_results = KnowledgeboxFindResults(
         resources=find_resources,
@@ -242,65 +233,4 @@ async def _external_index_retrieval(
         nodes=None,
     )
 
-    return retrieval_results, incomplete_results, query_parser
-
-
-@dataclass
-class ScoredParagraph:
-    id: str
-    score: float
-
-
-async def query_parser_from_find_request(
-    kbid: str, item: FindRequest, *, generative_model: Optional[str] = None
-) -> tuple[QueryParser, RankFusionAlgorithm, Reranker]:
-    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)
-
-    hidden = await filter_hidden_resources(kbid, item.show_hidden)
-
-    # XXX this is becoming the new /find query parsing, this should be moved to
-    # a cleaner abstraction
-
-    parsed = await parse_find(kbid, item)
-
-    rank_fusion = get_rank_fusion(parsed.rank_fusion)
-    reranker = get_reranker(parsed.reranker)
-
-    query_parser = QueryParser(
-        kbid=kbid,
-        features=item.features,
-        query=item.query,
-        query_entities=item.query_entities,
-        filter_expression=item.filter_expression,
-        faceted=None,
-        sort=None,
-        top_k=item.top_k,
-        min_score=item.min_score,
-        old_filters=OldFilterParams(
-            label_filters=item.filters,
-            keyword_filters=item.keyword_filters,
-            range_creation_start=item.range_creation_start,
-            range_creation_end=item.range_creation_end,
-            range_modification_start=item.range_modification_start,
-            range_modification_end=item.range_modification_end,
-            fields=item.fields,
-            key_filters=item.resource_filters,
-        ),
-        user_vector=item.vector,
-        vectorset=item.vectorset,
-        with_duplicates=item.with_duplicates,
-        with_synonyms=item.with_synonyms,
-        autofilter=item.autofilter,
-        security=item.security,
-        generative_model=generative_model,
-        rephrase=item.rephrase,
-        rephrase_prompt=item.rephrase_prompt,
-        hidden=hidden,
-        rank_fusion=rank_fusion,
-        reranker=reranker,
-    )
-    return (query_parser, rank_fusion, reranker)
+    return retrieval_results, incomplete_results, parsed

nucliadb/search/search/find_merge.py

@@ -32,6 +32,7 @@ from nucliadb.search.search.hydrator import (
     text_block_to_find_paragraph,
 )
 from nucliadb.search.search.merge import merge_relations_results
+from nucliadb.search.search.query_parser.models import UnitRetrieval
 from nucliadb.search.search.rank_fusion import RankFusionAlgorithm
 from nucliadb.search.search.rerankers import (
     RerankableItem,
@@ -51,7 +52,6 @@ from nucliadb_models.search import (
 )
 from nucliadb_protos.nodereader_pb2 import (
     DocumentScored,
-    EntitiesSubgraphRequest,
     ParagraphResult,
     ParagraphSearchResponse,
     RelationSearchResponse,
@@ -72,13 +72,10 @@ FIND_FETCH_OPS_DISTRIBUTION = metrics.Histogram(
 async def build_find_response(
     search_responses: list[SearchResponse],
     *,
+    retrieval: UnitRetrieval,
     kbid: str,
     query: str,
     rephrased_query: Optional[str],
-    relation_subgraph_query: EntitiesSubgraphRequest,
-    top_k: int,
-    min_score_bm25: float,
-    min_score_semantic: float,
     rank_fusion_algorithm: RankFusionAlgorithm,
     reranker: Reranker,
     show: list[ResourceProperties] = [],
@@ -86,6 +83,15 @@ async def build_find_response(
     field_type_filter: list[FieldTypeName] = [],
     highlight: bool = False,
 ) -> KnowledgeboxFindResults:
+    # XXX: we shouldn't need a min score that we haven't used. Previous
+    # implementations got this value from the proto request (i.e., default to 0)
+    min_score_bm25 = 0.0
+    if retrieval.query.keyword is not None:
+        min_score_bm25 = retrieval.query.keyword.min_score
+    min_score_semantic = 0.0
+    if retrieval.query.semantic is not None:
+        min_score_semantic = retrieval.query.semantic.min_score
+
     # merge
     search_response = merge_shard_responses(search_responses)
 
@@ -112,7 +118,7 @@ async def build_find_response(
         assert reranker.window is not None, "Reranker definition must enforce this condition"
         text_blocks_page, next_page = cut_page(merged_text_blocks, reranker.window)
     else:
-        text_blocks_page, next_page = cut_page(merged_text_blocks, top_k)
+        text_blocks_page, next_page = cut_page(merged_text_blocks, retrieval.top_k)
 
     # hydrate and rerank
     resource_hydration_options = ResourceHydrationOptions(
@@ -130,11 +136,14 @@ async def build_find_response(
         text_block_hydration_options=text_block_hydration_options,
         reranker=reranker,
         reranking_options=reranking_options,
-        top_k=top_k,
+        top_k=retrieval.top_k,
     )
 
     # build relations graph
-    relations = await merge_relations_results([search_response.relation], relation_subgraph_query)
+    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)
 
     # compose response
     find_resources = compose_find_resources(text_blocks, resources)
@@ -150,7 +159,7 @@ async def build_find_response(
         relations=relations,
         total=total_paragraphs,
         page_number=0,  # Bw/c with pagination
-        page_size=top_k,
+        page_size=retrieval.top_k,
         next_page=next_page,
         min_score=MinScore(bm25=_round(min_score_bm25), semantic=_round(min_score_semantic)),
     )

nucliadb/search/search/graph_strategy.py

@@ -16,7 +16,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/>.
-
 import heapq
 import json
 from collections import defaultdict
@@ -38,14 +37,16 @@ from nucliadb.search.search.chat.query import (
     find_request_from_ask_request,
     get_relations_results_from_entities,
 )
-from nucliadb.search.search.find import query_parser_from_find_request
 from nucliadb.search.search.find_merge import (
     compose_find_resources,
     hydrate_and_rerank,
 )
 from nucliadb.search.search.hydrator import ResourceHydrationOptions, TextBlockHydrationOptions
 from nucliadb.search.search.metrics import RAGMetrics
-from nucliadb.search.search.rerankers import Reranker, RerankingOptions
+from nucliadb.search.search.rerankers import (
+    Reranker,
+    RerankingOptions,
+)
 from nucliadb.search.utilities import get_predict
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.internal.predict import (
@@ -303,6 +304,7 @@ async def get_graph_results(
     user: str,
     origin: str,
     graph_strategy: GraphStrategy,
+    text_block_reranker: Reranker,
     generative_model: Optional[str] = None,
     metrics: RAGMetrics = RAGMetrics(),
     shards: Optional[list[str]] = None,
@@ -419,19 +421,16 @@ async def get_graph_results(
     # Get the text blocks of the paragraphs that contain the top relations
     with metrics.time("graph_strat_build_response"):
         find_request = find_request_from_ask_request(item, query)
-        query_parser, rank_fusion, reranker = await query_parser_from_find_request(
-            kbid, find_request, generative_model=generative_model
-        )
         find_results = await build_graph_response(
             kbid=kbid,
             query=query,
             final_relations=relations,
             scores=scores,
             top_k=graph_strategy.top_k,
-            reranker=reranker,
-            show=find_request.show,
-            extracted=find_request.extracted,
-            field_type_filter=find_request.field_type_filter,
+            reranker=text_block_reranker,
+            show=item.show,
+            extracted=item.extracted,
+            field_type_filter=item.field_type_filter,
             relation_text_as_paragraphs=graph_strategy.relation_text_as_paragraphs,
         )
     return find_results, find_request