nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/search/find_merge.py

@@ -18,37 +18,48 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-from typing import Any, Iterator, Optional, cast
+from typing import Iterable, Union
 
-from nucliadb_protos.nodereader_pb2 import (
-    DocumentScored,
-    EntitiesSubgraphRequest,
-    ParagraphResult,
-    SearchResponse,
-)
-
-from nucliadb.common.maindb.driver import Transaction
-from nucliadb.ingest.serialize import managed_serialize
-from nucliadb.middleware.transaction import get_read_only_transaction
+from nucliadb.common.external_index_providers.base import TextBlockMatch
+from nucliadb.common.ids import ParagraphId, VectorId
 from nucliadb.search import SERVICE_NAME, logger
-from nucliadb.search.search.cache import get_resource_cache
+from nucliadb.search.search.cut import cut_page
+from nucliadb.search.search.hydrator import (
+    ResourceHydrationOptions,
+    TextBlockHydrationOptions,
+    hydrate_resource_metadata,
+    hydrate_text_block,
+    text_block_to_find_paragraph,
+)
 from nucliadb.search.search.merge import merge_relations_results
+from nucliadb.search.search.rank_fusion import RankFusionAlgorithm
+from nucliadb.search.search.rerankers import (
+    RerankableItem,
+    Reranker,
+    RerankingOptions,
+)
 from nucliadb_models.common import FieldTypeName
-from nucliadb_models.resource import ExtractedDataTypeName
+from nucliadb_models.resource import ExtractedDataTypeName, Resource
 from nucliadb_models.search import (
     SCORE_TYPE,
     FindField,
-    FindParagraph,
     FindResource,
     KnowledgeboxFindResults,
     MinScore,
     ResourceProperties,
-    TempFindParagraph,
     TextPosition,
 )
+from nucliadb_protos.nodereader_pb2 import (
+    DocumentScored,
+    EntitiesSubgraphRequest,
+    ParagraphResult,
+    ParagraphSearchResponse,
+    RelationSearchResponse,
+    SearchResponse,
+    VectorSearchResponse,
+)
 from nucliadb_telemetry import metrics
 
-from . import paragraphs
 from .metrics import merge_observer
 
 FIND_FETCH_OPS_DISTRIBUTION = metrics.Histogram(
@@ -57,407 +68,413 @@ FIND_FETCH_OPS_DISTRIBUTION = metrics.Histogram(
 )
 
 
-def _round(x: float) -> float:
-    return round(x, ndigits=3)
-
-
-@merge_observer.wrap({"type": "set_text_value"})
-async def set_text_value(
+@merge_observer.wrap({"type": "find_merge"})
+async def build_find_response(
+    search_responses: list[SearchResponse],
+    *,
     kbid: str,
-    result_paragraph: TempFindParagraph,
-    max_operations: asyncio.Semaphore,
+    query: 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] = [],
+    extracted: list[ExtractedDataTypeName] = [],
+    field_type_filter: list[FieldTypeName] = [],
     highlight: bool = False,
-    ematches: Optional[list[str]] = None,
-    extracted_text_cache: Optional[paragraphs.ExtractedTextCache] = None,
-):
-    async with max_operations:
-        assert result_paragraph.paragraph
-        assert result_paragraph.paragraph.position
-        result_paragraph.paragraph.text = await paragraphs.get_paragraph_text(
-            kbid=kbid,
-            rid=result_paragraph.rid,
-            field=result_paragraph.field,
-            start=result_paragraph.paragraph.position.start,
-            end=result_paragraph.paragraph.position.end,
-            split=result_paragraph.split,
-            highlight=highlight,
-            ematches=ematches,
-            matches=[],  # TODO
-            extracted_text_cache=extracted_text_cache,
+) -> KnowledgeboxFindResults:
+    # merge
+    search_response = merge_shard_responses(search_responses)
+
+    keyword_results = keyword_results_to_text_block_matches(search_response.paragraph.results)
+    semantic_results = semantic_results_to_text_block_matches(
+        filter(
+            lambda x: x.score >= min_score_semantic,
+            search_response.vector.documents,
         )
+    )
+
+    merged_text_blocks: list[TextBlockMatch] = rank_fusion_algorithm.fuse(
+        keyword_results, semantic_results
+    )
+
+    # cut
+    # we assume pagination + predict reranker is forbidden and has been already
+    # enforced/validated by the query parsing.
+    if reranker.needs_extra_results:
+        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)
+
+    # hydrate and rerank
+    resource_hydration_options = ResourceHydrationOptions(
+        show=show, extracted=extracted, field_type_filter=field_type_filter
+    )
+    text_block_hydration_options = TextBlockHydrationOptions(
+        highlight=highlight,
+        ematches=search_response.paragraph.ematches,  # type: ignore
+    )
+    reranking_options = RerankingOptions(kbid=kbid, query=query)
+    text_blocks, resources, best_matches = await hydrate_and_rerank(
+        text_blocks_page,
+        kbid,
+        resource_hydration_options=resource_hydration_options,
+        text_block_hydration_options=text_block_hydration_options,
+        reranker=reranker,
+        reranking_options=reranking_options,
+        top_k=top_k,
+    )
+
+    # build relations graph
+    relations = await merge_relations_results([search_response.relation], relation_subgraph_query)
+
+    # compose response
+    find_resources = compose_find_resources(text_blocks, resources)
+
+    next_page = search_response.paragraph.next_page or next_page
+    total_paragraphs = search_response.paragraph.total
+
+    find_results = KnowledgeboxFindResults(
+        query=query,
+        resources=find_resources,
+        best_matches=best_matches,
+        relations=relations,
+        total=total_paragraphs,
+        page_number=0,  # Bw/c with pagination
+        page_size=top_k,
+        next_page=next_page,
+        min_score=MinScore(bm25=_round(min_score_bm25), semantic=_round(min_score_semantic)),
+    )
+    return find_results
+
+
+def merge_shard_responses(
+    responses: list[SearchResponse],
+) -> SearchResponse:
+    """Merge search responses into a single response as if there were no shards
+    involved.
+
+    ATENTION! This is not a complete merge, we are only merging the fields
+    needed to compose a /find response.
+
+    """
+    paragraphs = []
+    vectors = []
+    relations = []
+    for response in responses:
+        paragraphs.append(response.paragraph)
+        vectors.append(response.vector)
+        relations.append(response.relation)
 
-
-@merge_observer.wrap({"type": "set_resource_metadada_value"})
-async def set_resource_metadata_value(
-    txn: Transaction,
-    kbid: str,
-    resource: str,
-    show: list[ResourceProperties],
-    field_type_filter: list[FieldTypeName],
-    extracted: list[ExtractedDataTypeName],
-    find_resources: dict[str, FindResource],
-    max_operations: asyncio.Semaphore,
-):
-    async with max_operations:
-        serialized_resource = await managed_serialize(
-            txn,
-            kbid,
-            resource,
-            show,
-            field_type_filter=field_type_filter,
-            extracted=extracted,
-            service_name=SERVICE_NAME,
-        )
-        if serialized_resource is not None:
-            find_resources[resource].updated_from(serialized_resource)
-        else:
-            logger.warning(f"Resource {resource} not found in {kbid}")
-            find_resources.pop(resource, None)
-
-
-class Orderer:
-    def __init__(self):
-        self.boosted_items = []
-        self.items = []
-
-    def add(self, key: Any):
-        self.items.append(key)
-
-    def add_boosted(self, key: Any):
-        self.boosted_items.append(key)
-
-    def sorted_by_score(self) -> Iterator[Any]:
-        for key in sorted(self.items, key=lambda value: value[3], reverse=True):
-            yield key
-
-    def sorted_by_insertion(self) -> Iterator[Any]:
-        returned = set()
-        for key in self.boosted_items:
-            if key in returned:
-                continue
-            returned.add(key)
-            yield key
-
-        for key in self.items:
-            if key in returned:
-                continue
-            returned.add(key)
-            yield key
-
-
-@merge_observer.wrap({"type": "fetch_find_metadata"})
-async def fetch_find_metadata(
-    find_resources: dict[str, FindResource],
-    best_matches: list[str],
-    result_paragraphs: list[TempFindParagraph],
+    merged = SearchResponse(
+        paragraph=merge_shards_keyword_responses(paragraphs),
+        vector=merge_shards_semantic_responses(vectors),
+        relation=merge_shards_relation_responses(relations),
+    )
+    return merged
+
+
+def merge_shards_keyword_responses(
+    keyword_responses: list[ParagraphSearchResponse],
+) -> ParagraphSearchResponse:
+    """Merge keyword (paragraph) search responses into a single response as if
+    there were no shards involved.
+
+    ATENTION! This is not a complete merge, we are only merging the fields
+    needed to compose a /find response.
+
+    """
+    merged = ParagraphSearchResponse()
+    for response in keyword_responses:
+        merged.query = response.query
+        merged.next_page = merged.next_page or response.next_page
+        merged.total += response.total
+        merged.results.extend(response.results)
+        merged.ematches.extend(response.ematches)
+
+    return merged
+
+
+def merge_shards_semantic_responses(
+    semantic_responses: list[VectorSearchResponse],
+) -> VectorSearchResponse:
+    """Merge semantic (vector) search responses into a single response as if
+    there were no shards involved.
+
+    ATENTION! This is not a complete merge, we are only merging the fields
+    needed to compose a /find response.
+
+    """
+    merged = VectorSearchResponse()
+    for response in semantic_responses:
+        merged.documents.extend(response.documents)
+
+    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)
+
+    return merged
+
+
+def keyword_result_to_text_block_match(item: ParagraphResult) -> TextBlockMatch:
+    fuzzy_result = len(item.matches) > 0
+    return TextBlockMatch(
+        paragraph_id=ParagraphId.from_string(item.paragraph),
+        score=item.score.bm25,
+        score_type=SCORE_TYPE.BM25,
+        order=0,  # NOTE: this will be filled later
+        text="",  # NOTE: this will be filled later too
+        position=TextPosition(
+            page_number=item.metadata.position.page_number,
+            index=item.metadata.position.index,
+            start=item.start,
+            end=item.end,
+            start_seconds=[x for x in item.metadata.position.start_seconds],
+            end_seconds=[x for x in item.metadata.position.end_seconds],
+        ),
+        # XXX: we should split labels
+        field_labels=[],
+        paragraph_labels=list(item.labels),
+        fuzzy_search=fuzzy_result,
+        is_a_table=item.metadata.representation.is_a_table,
+        representation_file=item.metadata.representation.file,
+        page_with_visual=item.metadata.page_with_visual,
+    )
+
+
+def keyword_results_to_text_block_matches(items: Iterable[ParagraphResult]) -> list[TextBlockMatch]:
+    return [keyword_result_to_text_block_match(item) for item in items]
+
+
+class InvalidDocId(Exception):
+    """Raised while parsing an invalid id coming from semantic search"""
+
+    def __init__(self, invalid_vector_id: str):
+        self.invalid_vector_id = invalid_vector_id
+        super().__init__(f"Invalid vector ID: {invalid_vector_id}")
+
+
+def semantic_result_to_text_block_match(item: DocumentScored) -> TextBlockMatch:
+    try:
+        vector_id = VectorId.from_string(item.doc_id.id)
+    except (IndexError, ValueError):
+        raise InvalidDocId(item.doc_id.id)
+
+    return TextBlockMatch(
+        paragraph_id=ParagraphId.from_vector_id(vector_id),
+        score=item.score,
+        score_type=SCORE_TYPE.VECTOR,
+        order=0,  # NOTE: this will be filled later
+        text="",  # NOTE: this will be filled later too
+        position=TextPosition(
+            page_number=item.metadata.position.page_number,
+            index=item.metadata.position.index,
+            start=vector_id.vector_start,
+            end=vector_id.vector_end,
+            start_seconds=[x for x in item.metadata.position.start_seconds],
+            end_seconds=[x for x in item.metadata.position.end_seconds],
+        ),
+        # XXX: we should split labels
+        field_labels=[],
+        paragraph_labels=list(item.labels),
+        fuzzy_search=False,  # semantic search doesn't have fuzziness
+        is_a_table=item.metadata.representation.is_a_table,
+        representation_file=item.metadata.representation.file,
+        page_with_visual=item.metadata.page_with_visual,
+    )
+
+
+def semantic_results_to_text_block_matches(items: Iterable[DocumentScored]) -> list[TextBlockMatch]:
+    text_blocks: list[TextBlockMatch] = []
+    for item in items:
+        try:
+            text_block = semantic_result_to_text_block_match(item)
+        except InvalidDocId as exc:
+            logger.warning(f"Skipping invalid doc_id: {exc.invalid_vector_id}")
+            continue
+        text_blocks.append(text_block)
+    return text_blocks
+
+
+@merge_observer.wrap({"type": "hydrate_and_rerank"})
+async def hydrate_and_rerank(
+    text_blocks: Iterable[TextBlockMatch],
     kbid: str,
-    show: list[ResourceProperties],
-    field_type_filter: list[FieldTypeName],
-    extracted: list[ExtractedDataTypeName],
-    highlight: bool = False,
-    ematches: Optional[list[str]] = None,
-):
-    txn = await get_read_only_transaction()
-    resources = set()
-    operations = []
+    *,
+    resource_hydration_options: ResourceHydrationOptions,
+    text_block_hydration_options: TextBlockHydrationOptions,
+    reranker: Reranker,
+    reranking_options: RerankingOptions,
+    top_k: int,
+) -> tuple[list[TextBlockMatch], list[Resource], list[str]]:
+    """Given a list of text blocks from a retrieval operation, hydrate and
+    rerank the results.
+
+    This function returns either the entire list or a subset of updated
+    (hydrated and reranked) text blocks and their corresponding resource
+    metadata. It also returns an ordered list of best matches.
+
+    """
     max_operations = asyncio.Semaphore(50)
-    orderer = Orderer()
-    etcache = paragraphs.ExtractedTextCache()
-    for result_paragraph in result_paragraphs:
-        if result_paragraph.paragraph is not None:
-            find_resource = find_resources.setdefault(
-                result_paragraph.rid, FindResource(id=result_paragraph.id, fields={})
-            )
-            find_field = find_resource.fields.setdefault(
-                result_paragraph.field, FindField(paragraphs={})
-            )
 
-            if result_paragraph.paragraph.id in find_field.paragraphs:
-                # Its a multiple match, push the score
-                # find_field.paragraphs[result_paragraph.paragraph.id].score = 25
-                if (
-                    find_field.paragraphs[result_paragraph.paragraph.id].score
-                    < result_paragraph.paragraph.score
-                ):
-                    # Use Vector score if there are both
-                    find_field.paragraphs[result_paragraph.paragraph.id].score = (
-                        result_paragraph.paragraph.score * 2
-                    )
-                    orderer.add(
-                        (
-                            result_paragraph.rid,
-                            result_paragraph.field,
-                            result_paragraph.paragraph.id,
-                            result_paragraph.paragraph.score,
-                        )
-                    )
-                find_field.paragraphs[
-                    result_paragraph.paragraph.id
-                ].score_type = SCORE_TYPE.BOTH
-
-            else:
-                find_field.paragraphs[
-                    result_paragraph.paragraph.id
-                ] = result_paragraph.paragraph
-                orderer.add(
-                    (
-                        result_paragraph.rid,
-                        result_paragraph.field,
-                        result_paragraph.paragraph.id,
-                        result_paragraph.paragraph.score,
+    # Iterate text blocks and create text block and resource metadata hydration
+    # tasks depending on the reranker
+    text_blocks_by_id: dict[str, TextBlockMatch] = {}  # useful for faster access to text blocks later
+    resource_hydration_ops = {}
+    text_block_hydration_ops = []
+    for text_block in text_blocks:
+        rid = text_block.paragraph_id.rid
+        paragraph_id = text_block.paragraph_id.full()
+
+        # If we find multiple results (from different indexes) with different
+        # metadata, this statement will only get the metadata from the first on
+        # the list. We assume metadata is the same on all indexes, otherwise
+        # this would be a BUG
+        text_blocks_by_id.setdefault(paragraph_id, text_block)
+
+        # rerankers that need extra results may end with less resources than the
+        # ones we see now, so we'll skip this step and recompute the resources
+        # later
+        if not reranker.needs_extra_results:
+            if rid not in resource_hydration_ops:
+                resource_hydration_ops[rid] = asyncio.create_task(
+                    hydrate_resource_metadata(
+                        kbid,
+                        rid,
+                        options=resource_hydration_options,
+                        concurrency_control=max_operations,
+                        service_name=SERVICE_NAME,
                     )
                 )
 
-            operations.append(
-                asyncio.create_task(
-                    set_text_value(
-                        kbid=kbid,
-                        result_paragraph=result_paragraph,
-                        highlight=highlight,
-                        ematches=ematches,
-                        max_operations=max_operations,
-                        extracted_text_cache=etcache,
-                    )
-                )
-            )
-            resources.add(result_paragraph.rid)
-    etcache.clear()
-
-    for order, (rid, field_id, paragraph_id, _) in enumerate(orderer.sorted_by_score()):
-        find_resources[rid].fields[field_id].paragraphs[paragraph_id].order = order
-        best_matches.append(paragraph_id)
-
-    for resource in resources:
-        operations.append(
+        text_block_hydration_ops.append(
             asyncio.create_task(
-                set_resource_metadata_value(
-                    txn,
-                    kbid=kbid,
-                    resource=resource,
-                    show=show,
-                    field_type_filter=field_type_filter,
-                    extracted=extracted,
-                    find_resources=find_resources,
-                    max_operations=max_operations,
+                hydrate_text_block(
+                    kbid,
+                    text_block,
+                    text_block_hydration_options,
+                    concurrency_control=max_operations,
                 )
             )
         )
 
-    FIND_FETCH_OPS_DISTRIBUTION.observe(len(operations))
-    if len(operations) > 0:
-        done, _ = await asyncio.wait(operations)  # type: ignore
-        for task in done:
-            if task.exception() is not None:  # pragma: no cover
-                logger.error("Error fetching find metadata", exc_info=task.exception())
-
-
-@merge_observer.wrap({"type": "merge_paragraphs_vectors"})
-def merge_paragraphs_vectors(
-    paragraphs_shards: list[list[ParagraphResult]],
-    vectors_shards: list[list[DocumentScored]],
-    count: int,
-    page: int,
-    min_score: float,
-) -> tuple[list[TempFindParagraph], bool]:
-    merged_paragrahs: list[TempFindParagraph] = []
-
-    # We assume that paragraphs_shards and vectors_shards are already ordered
-    for paragraphs_shard in paragraphs_shards:
-        for paragraph in paragraphs_shard:
-            fuzzy_result = len(paragraph.matches) > 0
-            merged_paragrahs.append(
-                TempFindParagraph(
-                    paragraph_index=paragraph,
-                    field=paragraph.field,
-                    rid=paragraph.uuid,
-                    score=paragraph.score.bm25,
-                    start=paragraph.start,
-                    split=paragraph.split,
-                    end=paragraph.end,
-                    id=paragraph.paragraph,
-                    fuzzy_result=fuzzy_result,
-                    page_with_visual=paragraph.metadata.page_with_visual,
-                    reference=paragraph.metadata.representation.file,
-                    is_a_table=paragraph.metadata.representation.is_a_table,
+    # hydrate only the strictly needed before rerank
+    hydrated_text_blocks: list[TextBlockMatch]
+    hydrated_resources: list[Union[Resource, None]]
+
+    ops = [
+        *text_block_hydration_ops,
+        *resource_hydration_ops.values(),
+    ]
+    FIND_FETCH_OPS_DISTRIBUTION.observe(len(ops))
+    results = await asyncio.gather(*ops)
+
+    hydrated_text_blocks = results[: len(text_block_hydration_ops)]  # type: ignore
+    hydrated_resources = results[len(text_block_hydration_ops) :]  # type: ignore
+
+    # with the hydrated text, rerank and apply new scores to the text blocks
+    to_rerank = [
+        RerankableItem(
+            id=text_block.paragraph_id.full(),
+            score=text_block.score,
+            score_type=text_block.score_type,
+            content=text_block.text or "",  # TODO: add a warning, this shouldn't usually happen
+        )
+        for text_block in hydrated_text_blocks
+    ]
+    reranked = await reranker.rerank(to_rerank, reranking_options)
+
+    # after reranking, we can cut to the number of results the user wants, so we
+    # don't hydrate unnecessary stuff
+    reranked = reranked[:top_k]
+
+    matches = []
+    for item in reranked:
+        paragraph_id = item.id
+        score = item.score
+        score_type = item.score_type
+
+        text_block = text_blocks_by_id[paragraph_id]
+        text_block.score = score
+        text_block.score_type = score_type
+
+        matches.append((paragraph_id, score))
+
+    matches.sort(key=lambda x: x[1], reverse=True)
+
+    best_matches = []
+    best_text_blocks = []
+    resource_hydration_ops = {}
+    for order, (paragraph_id, _) in enumerate(matches):
+        text_block = text_blocks_by_id[paragraph_id]
+        text_block.order = order
+        best_matches.append(paragraph_id)
+        best_text_blocks.append(text_block)
+
+        # now we have removed the text block surplus, fetch resource metadata
+        if reranker.needs_extra_results:
+            rid = ParagraphId.from_string(paragraph_id).rid
+            if rid not in resource_hydration_ops:
+                resource_hydration_ops[rid] = asyncio.create_task(
+                    hydrate_resource_metadata(
+                        kbid,
+                        rid,
+                        options=resource_hydration_options,
+                        concurrency_control=max_operations,
+                        service_name=SERVICE_NAME,
+                    )
                 )
-            )
 
-    # merged_paragrahs.sort(key=lambda r: r.score, reverse=True)
+    # Finally, fetch resource metadata if we haven't already done it
+    if reranker.needs_extra_results:
+        ops = list(resource_hydration_ops.values())
+        FIND_FETCH_OPS_DISTRIBUTION.observe(len(ops))
+        hydrated_resources = await asyncio.gather(*ops)  # type: ignore
 
-    nextpos = 1
-    for vectors_shard in vectors_shards:
-        for vector in vectors_shard:
-            if vector.score < min_score:
-                logger.warning(
-                    f"Skipping low score vector: {vector.doc_id.id}. This should not happen"
-                )
-                continue
-            doc_id_split = vector.doc_id.id.split("/")
-            split = None
-            if len(doc_id_split) == 5:
-                rid, field_type, field, index, position = doc_id_split
-                paragraph_id = f"{rid}/{field_type}/{field}/{position}"
-            elif len(doc_id_split) == 6:
-                rid, field_type, field, split, index, position = doc_id_split
-                paragraph_id = f"{rid}/{field_type}/{field}/{split}/{position}"
-            else:
-                logger.warning(f"Skipping invalid doc_id: {vector.doc_id.id}")
-                continue
-            start, end = position.split("-")
-            merged_paragrahs.insert(
-                nextpos,
-                TempFindParagraph(
-                    vector_index=vector,
-                    rid=rid,
-                    field=f"/{field_type}/{field}",
-                    score=vector.score,
-                    start=int(start),
-                    end=int(end),
-                    split=split,
-                    id=paragraph_id,
-                ),
-            )
-            nextpos += 3
-
-    # merged_paragrahs.sort(key=lambda r: r.score, reverse=True)
-    init_position = count * page
-    end_position = init_position + count
-    next_page = len(merged_paragrahs) > end_position
-    merged_paragrahs = merged_paragrahs[init_position:end_position]
-
-    for merged_paragraph in merged_paragrahs:
-        if merged_paragraph.vector_index is not None:
-            merged_paragraph.paragraph = FindParagraph(
-                score=merged_paragraph.vector_index.score,
-                score_type=SCORE_TYPE.VECTOR,
-                text="",
-                labels=[],  # TODO: Get labels from index
-                page_with_visual=merged_paragraph.vector_index.metadata.page_with_visual,
-                reference=merged_paragraph.vector_index.metadata.representation.file,
-                is_a_table=merged_paragraph.vector_index.metadata.representation.is_a_table,
-                position=TextPosition(
-                    page_number=merged_paragraph.vector_index.metadata.position.page_number,
-                    index=merged_paragraph.vector_index.metadata.position.index,
-                    start=merged_paragraph.start,
-                    end=merged_paragraph.end,
-                    start_seconds=[
-                        x
-                        for x in merged_paragraph.vector_index.metadata.position.start_seconds
-                    ],
-                    end_seconds=[
-                        x
-                        for x in merged_paragraph.vector_index.metadata.position.end_seconds
-                    ],
-                ),
-                id=merged_paragraph.id,
-                # Vector searches don't have fuzziness
-                fuzzy_result=False,
-            )
-        elif merged_paragraph.paragraph_index is not None:
-            merged_paragraph.paragraph = FindParagraph(
-                score=merged_paragraph.paragraph_index.score.bm25,
-                score_type=SCORE_TYPE.BM25,
-                text="",
-                labels=[x for x in merged_paragraph.paragraph_index.labels],
-                page_with_visual=merged_paragraph.paragraph_index.metadata.page_with_visual,
-                reference=merged_paragraph.paragraph_index.metadata.representation.file,
-                is_a_table=merged_paragraph.paragraph_index.metadata.representation.is_a_table,
-                position=TextPosition(
-                    page_number=merged_paragraph.paragraph_index.metadata.position.page_number,
-                    index=merged_paragraph.paragraph_index.metadata.position.index,
-                    start=merged_paragraph.start,
-                    end=merged_paragraph.end,
-                    start_seconds=[
-                        x
-                        for x in merged_paragraph.paragraph_index.metadata.position.start_seconds
-                    ],
-                    end_seconds=[
-                        x
-                        for x in merged_paragraph.paragraph_index.metadata.position.end_seconds
-                    ],
-                ),
-                id=merged_paragraph.id,
-                fuzzy_result=merged_paragraph.fuzzy_result,
-            )
-    return merged_paragrahs, next_page
+    resources = [resource for resource in hydrated_resources if resource is not None]
 
+    return best_text_blocks, resources, best_matches
 
-@merge_observer.wrap({"type": "find_merge"})
-async def find_merge_results(
-    search_responses: list[SearchResponse],
-    count: int,
-    page: int,
-    kbid: str,
-    show: list[ResourceProperties],
-    field_type_filter: list[FieldTypeName],
-    extracted: list[ExtractedDataTypeName],
-    requested_relations: EntitiesSubgraphRequest,
-    min_score_bm25: float,
-    min_score_semantic: float,
-    highlight: bool = False,
-) -> KnowledgeboxFindResults:
-    # force getting transaction on current asyncio task
-    # so all sub tasks will use the same transaction
-    # this is contextvar magic that is probably not ideal
-    await get_read_only_transaction()
 
-    paragraphs: list[list[ParagraphResult]] = []
-    vectors: list[list[DocumentScored]] = []
-    relations = []
+def compose_find_resources(
+    text_blocks: list[TextBlockMatch],
+    resources: list[Resource],
+) -> dict[str, FindResource]:
+    find_resources: dict[str, FindResource] = {}
 
-    next_page = True
-    ematches: list[str] = []
-    real_query = ""
-    total_paragraphs = 0
-    for response in search_responses:
-        # Iterate over answers from different logic shards
+    for resource in resources:
+        rid = resource.id
+        if rid not in find_resources:
+            find_resources[rid] = FindResource(id=rid, fields={})
+            find_resources[rid].updated_from(resource)
 
-        ematches.extend(response.paragraph.ematches)
-        real_query = response.paragraph.query
-        next_page = next_page and response.paragraph.next_page
-        total_paragraphs += response.paragraph.total
+    for text_block in text_blocks:
+        rid = text_block.paragraph_id.rid
+        if rid not in find_resources:
+            # resource not found in db, skipping
+            continue
 
-        paragraphs.append(cast(list[ParagraphResult], response.paragraph.results))
-        vectors.append(cast(list[DocumentScored], response.vector.documents))
+        find_resource = find_resources[rid]
+        field_id = text_block.paragraph_id.field_id.short_without_subfield()
+        find_field = find_resource.fields.setdefault(field_id, FindField(paragraphs={}))
 
-        relations.append(response.relation)
+        paragraph_id = text_block.paragraph_id.full()
+        find_paragraph = text_block_to_find_paragraph(text_block)
 
-    rcache = get_resource_cache(clear=True)
+        find_field.paragraphs[paragraph_id] = find_paragraph
 
-    try:
-        result_paragraphs, merged_next_page = merge_paragraphs_vectors(
-            paragraphs, vectors, count, page, min_score_semantic
-        )
-        next_page = next_page or merged_next_page
-
-        api_results = KnowledgeboxFindResults(
-            resources={},
-            query=real_query,
-            total=total_paragraphs,
-            page_number=page,
-            page_size=count,
-            next_page=next_page,
-            min_score=MinScore(
-                bm25=_round(min_score_bm25), semantic=_round(min_score_semantic)
-            ),
-            best_matches=[],
-        )
+    return find_resources
 
-        await fetch_find_metadata(
-            api_results.resources,
-            api_results.best_matches,
-            result_paragraphs,
-            kbid,
-            show,
-            field_type_filter,
-            extracted,
-            highlight,
-            ematches,
-        )
-        api_results.relations = await merge_relations_results(
-            relations, requested_relations
-        )
 
-        return api_results
-    finally:
-        rcache.clear()
+def _round(x: float) -> float:
+    return round(x, ndigits=3)