hyperforge-nucliadb-agentic 1.0.0.post64__py3-none-any.whl

hyperforge_nucliadb_agentic/ask/search/graph_strategy.py

@@ -0,0 +1,1138 @@
+import heapq
+import json
+from collections import defaultdict
+from collections.abc import Collection, Iterable
+from dataclasses import dataclass
+from typing import Any
+
+from nuclia_models.predict.generative_responses import (
+    JSONGenerativeResponse,
+    MetaGenerativeResponse,
+    StatusGenerativeResponse,
+)
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.graph.requests import (
+    And,
+    AnyNode,
+    Generated,
+    Generator,
+    GraphNode,
+    GraphNodesSearchRequest,
+    GraphPath,
+    GraphPathQuery,
+    GraphSearchRequest,
+    NodeMatchKindName,
+    Not,
+    Or,
+    Relation,
+    RelationNodeType,
+    RelationType,
+)
+from nucliadb_models.graph.responses import GraphSearchResponse
+from nucliadb_models.internal.predict import RerankModel
+from nucliadb_models.metadata import RelationMetadata
+from nucliadb_models.resource import ExtractedDataTypeName
+from nucliadb_models.retrieval import GraphScore
+from nucliadb_models.search import (
+    SCORE_TYPE,
+    DirectionalRelation,
+    EntitySubgraph,
+    NucliaDBClientType,
+    RelatedEntities,
+    RelatedEntity,
+    RelationDirection,
+)
+from nucliadb_protos import utils_pb2
+from nucliadb_protos.utils_pb2 import RelationNode
+from nucliadb_sdk.v2 import NucliaDBAsync
+from pydantic import BaseModel
+from sentry_sdk import capture_exception
+
+from hyperforge_nucliadb_agentic.ask import logger
+from hyperforge_nucliadb_agentic.ask.model import (
+    AskRequest,
+    ChatModel,
+    FindRequest,
+    GraphStrategy,
+    KnowledgeboxFindResults,
+    QueryEntityDetection,
+    RelationRanking,
+    Relations,
+    ResourceProperties,
+    TextPosition,
+    UserPrompt,
+)
+from hyperforge_nucliadb_agentic.ask.predict import get_predict
+from hyperforge_nucliadb_agentic.ask.search import rpc
+from hyperforge_nucliadb_agentic.ask.search.hydrator import (
+    ResourceHydrationOptions,
+    TextBlockHydrationOptions,
+)
+from hyperforge_nucliadb_agentic.ask.search.metrics import Metrics
+from hyperforge_nucliadb_agentic.ask.search.rerankers import (
+    Reranker,
+    RerankingOptions,
+)
+from hyperforge_nucliadb_agentic.ask.search.retrieval import (
+    compose_find_resources,
+    find_request_from_ask_request,
+    hydrate_and_rerank,
+)
+from hyperforge_nucliadb_agentic.ask.utils.ids import (
+    FieldId,
+    ParagraphId,
+)
+from hyperforge_nucliadb_agentic.ask.utils.text_blocks import (
+    TextBlockMatch,
+)
+
+RelationNodeTypeMap: dict[
+    RelationNodeType, utils_pb2.RelationNode.NodeType.ValueType
+] = {
+    RelationNodeType.ENTITY: utils_pb2.RelationNode.NodeType.ENTITY,
+    RelationNodeType.LABEL: utils_pb2.RelationNode.NodeType.LABEL,
+    RelationNodeType.RESOURCE: utils_pb2.RelationNode.NodeType.RESOURCE,
+    RelationNodeType.USER: utils_pb2.RelationNode.NodeType.USER,
+}
+
+RelationNodeTypePbMap: dict[
+    utils_pb2.RelationNode.NodeType.ValueType, RelationNodeType
+] = {
+    utils_pb2.RelationNode.NodeType.ENTITY: RelationNodeType.ENTITY,
+    utils_pb2.RelationNode.NodeType.LABEL: RelationNodeType.LABEL,
+    utils_pb2.RelationNode.NodeType.RESOURCE: RelationNodeType.RESOURCE,
+    utils_pb2.RelationNode.NodeType.USER: RelationNodeType.USER,
+}
+
+SCHEMA = {
+    "title": "score_triplets",
+    "description": "Return a list of triplets and their relevance scores (0-10) for the supplied question.",
+    "type": "object",
+    "properties": {
+        "triplets": {
+            "type": "array",
+            "description": "A list of triplets with their relevance scores.",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "head_entity": {
+                        "type": "string",
+                        "description": "The first entity in the triplet.",
+                    },
+                    "relationship": {
+                        "type": "string",
+                        "description": "The relationship between the two entities.",
+                    },
+                    "tail_entity": {
+                        "type": "string",
+                        "description": "The second entity in the triplet.",
+                    },
+                    "score": {
+                        "type": "integer",
+                        "description": "A relevance score in the range 0 to 10.",
+                        "minimum": 0,
+                        "maximum": 10,
+                    },
+                },
+                "required": ["head_entity", "relationship", "tail_entity", "score"],
+            },
+        }
+    },
+    "required": ["triplets"],
+}
+
+PROMPT = """\
+You are an advanced language model assisting in scoring relationships (edges) between two entities in a knowledge graph, given a user's question.
+
+For each provided **(head_entity, relationship, tail_entity)**, you must:
+1. Assign a **relevance score** between **0** and **10**.
+2. **0** means “this relationship can't be relevant at all to the question.”
+3. **10** means “this relationship is extremely relevant to the question.”
+4. You may use **any integer** between 0 and 10 (e.g., 3, 7, etc.) based on how relevant you deem the relationship to be.
+5. **Language Agnosticism**: The question and the relationships may be in different languages. The relevance scoring should still work and be agnostic of the language.
+6. Relationships that may not answer the question directly but expand knowledge in a relevant way, should also be scored positively.
+
+Once you have decided the best score for each triplet, return these results **using a function call** in JSON format with the following rules:
+
+- The function name should be `score_triplets`.
+- The first argument should be the list of triplets.
+- Each triplet should have the following keys:
+  - `head_entity`: The first entity in the triplet.
+  - `relationship`: The relationship between the two entities.
+  - `tail_entity`: The second entity in the triplet.
+  - `score`: The relevance score in the range 0 to 10.
+
+You **must** comply with the provided JSON Schema to ensure a well-structured response and mantain the order of the triplets.
+
+
+## Examples:
+
+### Example 1:
+
+**Input**
+
+{
+ "question": "Who is the mayor of the capital city of Australia?",
+ "triplets": [
+    {
+        "head_entity": "Australia",
+        "relationship": "has prime minister",
+        "tail_entity": "Scott Morrison"
+    },
+    {
+        "head_entity": "Canberra",
+        "relationship": "is capital of",
+        "tail_entity": "Australia"
+    },
+    {
+        "head_entity": "Scott Knowles",
+        "relationship": "holds position",
+        "tail_entity": "Mayor"
+    },
+    {
+        "head_entity": "Barbera Smith",
+        "relationship": "tiene cargo",
+        "tail_entity": "Alcalde"
+    },
+    {
+        "head_entity": "Austria",
+        "relationship": "has capital",
+        "tail_entity": "Vienna"
+    }
+ ]
+}
+
+**Output**
+
+{
+    "triplets": [
+        {
+            "head_entity": "Australia",
+            "relationship": "has prime minister",
+            "tail_entity": "Scott Morrison",
+            "score": 4
+        },
+        {
+            "head_entity": "Canberra",
+            "relationship": "is capital of",
+            "tail_entity": "Australia",
+            "score": 8
+        },
+        {
+            "head_entity": "Scott Knowles",
+            "relationship": "holds position",
+            "tail_entity": "Mayor",
+            "score": 8
+        },
+        {
+            "head_entity": "Barbera Smith",
+            "relationship": "tiene cargo",
+            "tail_entity": "Alcalde",
+            "score": 8
+        },
+        {
+            "head_entity": "Austria",
+            "relationship": "has capital",
+            "tail_entity": "Vienna",
+            "score": 0
+        }
+    ]
+}
+
+
+
+### Example 2:
+
+**Input**
+
+{
+    "question": "How many products does John Adams Roofing Inc. offer?",
+    "triplets": [
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "has product",
+            "tail_entity": "Titanium Grade 3 Roofing Nails"
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "is located in",
+            "tail_entity": "New York"
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "was founded by",
+            "tail_entity": "John Adams"
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "tiene stock",
+            "tail_entity": "Baldosas solares"
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "has product",
+            "tail_entity": "Mercerized Cotton Thread"
+        }
+    ]
+}
+
+**Output**
+
+{
+    "triplets": [
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "has product",
+            "tail_entity": "Titanium Grade 3 Roofing Nails",
+            "score": 10
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "is located in",
+            "tail_entity": "New York",
+            "score": 6
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "was founded by",
+            "tail_entity": "John Adams",
+            "score": 5
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "tiene stock",
+            "tail_entity": "Baldosas solares",
+            "score": 10
+        },
+        {
+            "head_entity": "John Adams Roofing Inc.",
+            "relationship": "has product",
+            "tail_entity": "Mercerized Cotton Thread",
+            "score": 10
+        }
+    ]
+}
+
+Now, let's get started! Here are the triplets you need to score:
+
+**Input**
+
+"""
+
+
+@dataclass(frozen=True)
+class FrozenRelationNode:
+    ntype: RelationNode.NodeType.ValueType
+    subtype: str
+    value: str
+
+
+def freeze_node(r: RelationNode) -> FrozenRelationNode:
+    return FrozenRelationNode(ntype=r.ntype, subtype=r.subtype, value=r.value)
+
+
+class RelationsParagraphMatch(BaseModel):
+    paragraph_id: ParagraphId
+    score: float
+    relations: Relations
+
+
+async def get_graph_results(
+    *,
+    search_sdk: NucliaDBAsync,
+    kbid: str,
+    query: str,
+    item: AskRequest,
+    ndb_client: NucliaDBClientType,
+    user: str,
+    origin: str,
+    graph_strategy: GraphStrategy,
+    text_block_reranker: Reranker,
+    metrics: Metrics,
+    generative_model: str | None = None,
+    shards: list[str] | None = None,
+) -> tuple[KnowledgeboxFindResults, FindRequest]:
+    relations = Relations(entities={})
+    explored_entities: set[FrozenRelationNode] = set()
+    scores: dict[str, list[float]] = {}
+    predict = get_predict()
+    entities_to_explore: list[RelationNode] = []
+
+    for hop in range(graph_strategy.hops):
+        if hop == 0:
+            # Get the entities from the query
+            with metrics.time("graph_strat_query_entities"):
+                if (
+                    graph_strategy.query_entity_detection
+                    == QueryEntityDetection.SUGGEST
+                ):
+                    relation_result = await fuzzy_search_entities(
+                        search_sdk,
+                        kbid=kbid,
+                        query=query,
+                    )
+                    if relation_result is not None:
+                        entities_to_explore = [
+                            RelationNode(
+                                ntype=RelationNode.NodeType.ENTITY,
+                                value=result.value,
+                                subtype=result.family,
+                            )
+                            for result in relation_result.entities
+                        ]
+                elif (
+                    not entities_to_explore
+                    or graph_strategy.query_entity_detection
+                    == QueryEntityDetection.PREDICT
+                ):
+                    try:
+                        # Purposely ignore the entity subtype. This is done so we find all entities that match
+                        # the entity by name. e.g: in a query like "2000", predict might detect the number as
+                        # a year entity or as a currency entity. We want graph results for both, so we ignore the
+                        # subtype just in this case.
+                        entities_to_explore = [
+                            RelationNode(ntype=r.ntype, value=r.value, subtype="")
+                            for r in await predict.detect_entities(kbid, query)
+                        ]
+                    except Exception as e:
+                        capture_exception(e)
+                        logger.exception(
+                            "Error in detecting entities for graph strategy"
+                        )
+                        entities_to_explore = []
+        else:
+            # Find neighbors of the current relations and remove the ones already explored
+            entities_to_explore = [
+                RelationNode(
+                    ntype=RelationNode.NodeType.ENTITY,
+                    value=relation.entity,
+                    subtype=relation.entity_subtype,
+                )
+                for subgraph in relations.entities.values()
+                for relation in subgraph.related_to
+                if FrozenRelationNode(
+                    ntype=RelationNodeTypeMap[relation.entity_type],
+                    subtype=relation.entity_subtype,
+                    value=relation.entity,
+                )
+                not in explored_entities
+            ]
+
+        if not entities_to_explore:
+            break
+
+        # Get the relations for the new entities
+        with metrics.time("graph_strat_neighbor_relations"):
+            try:
+                graph_response = await find_graph_neighbours(
+                    search_sdk,
+                    kbid,
+                    entities_to_explore,
+                    explored_entities,
+                    exclude_processor_relations=graph_strategy.exclude_processor_relations,
+                )
+                new_relations = merge_relations_results(
+                    graph_response,
+                    entities_to_explore,
+                    only_with_metadata=not graph_strategy.relation_text_as_paragraphs,
+                )
+            except Exception as e:
+                capture_exception(e)
+                logger.exception("Error in getting query relations for graph strategy")
+                new_relations = Relations(entities={})
+
+            relations.entities.update(new_relations.entities)
+            discovered_entities = []
+
+            for path in graph_response.paths:
+                for path_node in [path.source, path.destination]:
+                    node = RelationNode(
+                        value=path_node.value,
+                        ntype=RelationNodeTypeMap[path_node.type],
+                        subtype=path_node.group,
+                    )
+                    if (
+                        node not in entities_to_explore
+                        and freeze_node(node) not in explored_entities
+                    ):
+                        discovered_entities.append(node)
+
+            if not discovered_entities:
+                break
+
+            explored_entities.update([freeze_node(n) for n in entities_to_explore])
+            entities_to_explore = discovered_entities
+
+        # Rank the relevance of the relations
+        with metrics.time("graph_strat_rank_relations"):
+            try:
+                if graph_strategy.relation_ranking == RelationRanking.RERANKER:
+                    relations, scores = await rank_relations_reranker(
+                        relations,
+                        query,
+                        kbid,
+                        user,
+                        top_k=graph_strategy.top_k,
+                    )
+                elif graph_strategy.relation_ranking == RelationRanking.GENERATIVE:
+                    relations, scores = await rank_relations_generative(
+                        relations,
+                        query,
+                        kbid,
+                        user,
+                        top_k=graph_strategy.top_k,
+                        generative_model=generative_model,
+                    )
+            except Exception as e:
+                capture_exception(e)
+                logger.exception("Error in ranking relations for graph strategy")
+                relations = Relations(entities={})
+                scores = {}
+                break
+
+    # 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)
+        find_results = await build_graph_response(
+            search_sdk=search_sdk,
+            kbid=kbid,
+            query=query,
+            final_relations=relations,
+            scores=scores,
+            top_k=graph_strategy.top_k,
+            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
+
+
+async def fuzzy_search_entities(
+    search_sdk: NucliaDBAsync,
+    kbid: str,
+    query: str,
+) -> RelatedEntities | None:
+    """Fuzzy find entities in KB given a query using the same methodology as /suggest, but split by words."""
+
+    # Build an OR for each word in the query matching with fuzzy any word in any
+    # node in any position. I.e., for the query "Rose Hamiltn", it'll match
+    # "Rosa Parks" and "Margaret Hamilton"
+    operands = []
+    for word in query.split():
+        if len(word) >= 3:
+            subquery = AnyNode(value=word, match=NodeMatchKindName.FUZZY_WORDS)
+        else:
+            # words smaller than 3 characters are too noisy and thus, forbidden in the API
+            subquery = AnyNode(value=word, match=NodeMatchKindName.EXACT)
+        operands.append(subquery)
+
+    if len(operands) == 0:
+        return RelatedEntities(entities=[], total=0)
+
+    # XXX: are those enough results? Too many?
+    top_k = 50
+    if len(operands) == 1:
+        request = GraphNodesSearchRequest(query=operands[0], top_k=top_k)
+    else:
+        request = GraphNodesSearchRequest(query=Or(operands=operands), top_k=top_k)
+
+    try:
+        response = await rpc.graph_nodes(search_sdk, kbid, request)
+    except Exception as exc:
+        capture_exception(exc)
+        logger.exception("Error in finding entities in query for graph strategy")
+        return None
+
+    # merge shard results while deduplicating repeated entities across shards
+    unique_entities: set[RelatedEntity] = {
+        RelatedEntity(family=e.group, value=e.value) for e in response.nodes
+    }
+
+    return RelatedEntities(entities=list(unique_entities), total=len(unique_entities))
+
+
+async def rank_relations_reranker(
+    relations: Relations,
+    query: str,
+    kbid: str,
+    user: str,
+    top_k: int,
+    score_threshold: float = 0.02,
+) -> tuple[Relations, dict[str, list[float]]]:
+    # Store the index for keeping track after scoring
+    flat_rels: list[tuple[str, int, DirectionalRelation]] = [
+        (ent, idx, rel)
+        for (ent, rels) in relations.entities.items()
+        for (idx, rel) in enumerate(rels.related_to)
+    ]
+    # Build triplets (dict) from each relation for use in reranker
+    triplets: list[dict[str, str]] = [
+        {
+            "head_entity": ent,
+            "relationship": rel.relation_label,
+            "tail_entity": rel.entity,
+        }
+        if rel.direction == RelationDirection.OUT
+        else {
+            "head_entity": rel.entity,
+            "relationship": rel.relation_label,
+            "tail_entity": ent,
+        }
+        for (ent, _, rel) in flat_rels
+    ]
+
+    # Dedupe triplets so that they get evaluated once; map triplet -> [orig_indices]
+    triplet_to_orig_indices: dict[tuple[str, str, str], list[int]] = {}
+    unique_triplets: list[dict[str, str]] = []
+
+    for i, t in enumerate(triplets):
+        key = (t["head_entity"], t["relationship"], t["tail_entity"])
+        if key not in triplet_to_orig_indices:
+            triplet_to_orig_indices[key] = []
+            unique_triplets.append(t)
+        triplet_to_orig_indices[key].append(i)
+
+    # Build the reranker model input
+    predict = get_predict()
+    rerank_model = RerankModel(
+        question=query,
+        user_id=user,
+        context={
+            str(idx): f"{t['head_entity']} {t['relationship']} {t['tail_entity']}"
+            for idx, t in enumerate(unique_triplets)
+        },
+    )
+    # Get the rerank scores
+    res = await predict.rerank(kbid, rerank_model)
+
+    # Convert returned scores to a list of (int_idx, score)
+    # where int_idx corresponds to indices in unique_triplets
+    reranked_indices_scores = [
+        (int(idx), score) for idx, score in res.context_scores.items()
+    ]
+
+    return _scores_to_ranked_rels(
+        unique_triplets,
+        reranked_indices_scores,
+        triplet_to_orig_indices,
+        flat_rels,
+        top_k,
+        score_threshold,
+    )
+
+
+async def rank_relations_generative(
+    relations: Relations,
+    query: str,
+    kbid: str,
+    user: str,
+    top_k: int,
+    generative_model: str | None = None,
+    score_threshold: float = 2,
+    max_rels_to_eval: int = 100,
+) -> tuple[Relations, dict[str, list[float]]]:
+    # Store the index for keeping track after scoring
+    flat_rels: list[tuple[str, int, DirectionalRelation]] = [
+        (ent, idx, rel)
+        for (ent, rels) in relations.entities.items()
+        for (idx, rel) in enumerate(rels.related_to)
+    ]
+    triplets: list[dict[str, str]] = [
+        {
+            "head_entity": ent,
+            "relationship": rel.relation_label,
+            "tail_entity": rel.entity,
+        }
+        if rel.direction == RelationDirection.OUT
+        else {
+            "head_entity": rel.entity,
+            "relationship": rel.relation_label,
+            "tail_entity": ent,
+        }
+        for (ent, _, rel) in flat_rels
+    ]
+
+    # Dedupe triplets so that they get evaluated once, we will re-associate the scores later
+    triplet_to_orig_indices: dict[tuple[str, str, str], list[int]] = {}
+    unique_triplets: list[dict[str, str]] = []
+
+    for i, t in enumerate(triplets):
+        key = (t["head_entity"], t["relationship"], t["tail_entity"])
+        if key not in triplet_to_orig_indices:
+            triplet_to_orig_indices[key] = []
+            unique_triplets.append(t)
+        triplet_to_orig_indices[key].append(i)
+
+    if len(flat_rels) > max_rels_to_eval:
+        logger.warning(
+            f"Too many relations to evaluate ({len(flat_rels)}), using reranker to reduce"
+        )
+        return await rank_relations_reranker(
+            relations, query, kbid, user, top_k=max_rels_to_eval
+        )
+
+    data = {
+        "question": query,
+        "triplets": unique_triplets,
+    }
+    prompt = PROMPT + json.dumps(data, indent=4)
+
+    predict = get_predict()
+    chat_model = ChatModel(
+        question=prompt,
+        user_id=user,
+        json_schema=SCHEMA,
+        format_prompt=False,  # We supply our own prompt
+        query_context_order={},
+        query_context={},
+        user_prompt=UserPrompt(prompt=prompt),
+        max_tokens=4096,
+        generative_model=generative_model,
+    )
+
+    ident, model, answer_stream = await predict.chat_query_ndjson(kbid, chat_model)
+    response_json = None
+    status = None
+    _ = None
+
+    async for generative_chunk in answer_stream:
+        item = generative_chunk.chunk
+        if isinstance(item, JSONGenerativeResponse):
+            response_json = item
+        elif isinstance(item, StatusGenerativeResponse):
+            status = item
+        elif isinstance(item, MetaGenerativeResponse):
+            _ = item
+        else:
+            raise ValueError(f"Unknown generative chunk type: {item}")
+
+    if response_json is None or status is None or status.code != "0":
+        raise ValueError("No JSON response found")
+
+    scored_unique_triplets: list[dict[str, str | Any]] = response_json.object[
+        "triplets"
+    ]
+
+    if len(scored_unique_triplets) != len(unique_triplets):
+        raise ValueError("Mismatch between input and output triplets")
+
+    unique_indices_scores = (
+        (idx, float(t["score"])) for (idx, t) in enumerate(scored_unique_triplets)
+    )
+
+    return _scores_to_ranked_rels(
+        unique_triplets,
+        unique_indices_scores,
+        triplet_to_orig_indices,
+        flat_rels,
+        top_k,
+        score_threshold,
+    )
+
+
+def _scores_to_ranked_rels(
+    unique_triplets: list[dict[str, str]],
+    unique_indices_scores: Iterable[tuple[int, float]],
+    triplet_to_orig_indices: dict[tuple[str, str, str], list[int]],
+    flat_rels: list[tuple[str, int, DirectionalRelation]],
+    top_k: int,
+    score_threshold: float,
+) -> tuple[Relations, dict[str, list[float]]]:
+    """
+    Helper function to convert unique scores assigned by a model back to the original relations while taking
+    care of threshold
+    """
+    top_k_indices_scores = heapq.nlargest(
+        top_k, unique_indices_scores, key=lambda x: x[1]
+    )
+
+    # Prepare a new Relations object + a dict of top scores by entity
+    top_k_rels: dict[str, EntitySubgraph] = defaultdict(
+        lambda: EntitySubgraph(related_to=[])
+    )
+    top_k_scores_by_ent: dict[str, list[float]] = defaultdict(list)
+    # Re-expand model scores to the original triplets
+    for idx, score in top_k_indices_scores:
+        # If the model's score is below threshold, skip
+        if score <= score_threshold:
+            continue
+
+        # Identify which original triplets (in flat_rels) this corresponds to
+        t = unique_triplets[idx]
+        key = (t["head_entity"], t["relationship"], t["tail_entity"])
+        orig_indices = triplet_to_orig_indices[key]
+
+        for orig_i in orig_indices:
+            ent, rel_idx, rel = flat_rels[orig_i]
+            # Insert the relation into top_k_rels
+            top_k_rels[ent].related_to.append(rel)
+
+            # Keep track of which indices were chosen per entity
+            top_k_scores_by_ent[ent].append(score)
+
+    return Relations(entities=top_k_rels), dict(top_k_scores_by_ent)
+
+
+def build_text_blocks_from_relations(
+    relations: Relations,
+    scores: dict[str, list[float]],
+) -> list[TextBlockMatch]:
+    """
+    The goal of this function is to generate  TextBlockMatch with custom text for each unique relation in the graph.
+
+    This is a hacky way to generate paragraphs from relations, and it is not the intended use of TextBlockMatch.
+    """
+    # Build a set of unique triplets with their scores
+    triplets: dict[
+        tuple[str, str, str], tuple[float, Relations, ParagraphId | None]
+    ] = defaultdict(lambda: (0.0, Relations(entities={}), None))
+    paragraph_count = 0
+    for ent, subgraph in relations.entities.items():
+        if ent not in scores:
+            continue
+        for rel, score in zip(subgraph.related_to, scores[ent]):
+            key = (
+                (
+                    ent,
+                    rel.relation_label,
+                    rel.entity,
+                )
+                if rel.direction == RelationDirection.OUT
+                else (rel.entity, rel.relation_label, ent)
+            )
+            existing_score, existing_relations, p_id = triplets[key]
+            if ent not in existing_relations.entities:
+                existing_relations.entities[ent] = EntitySubgraph(related_to=[])
+
+            # XXX: Since relations with the same triplet can point to different paragraphs,
+            # we keep the first one, but we lose the other ones
+            if p_id is None and rel.metadata and rel.metadata.paragraph_id:
+                p_id = ParagraphId.from_string(rel.metadata.paragraph_id)
+            else:
+                # No paragraph ID set, fake it so we can hydrate the resource
+                p_id = ParagraphId(
+                    field_id=FieldId(rel.resource_id, "a", "usermetadata"),
+                    paragraph_start=paragraph_count,
+                    paragraph_end=paragraph_count + 1,
+                )
+                paragraph_count += 1
+            existing_relations.entities[ent].related_to.append(rel)
+            # XXX: Here we use the max even though all relations with same triplet should have same score
+            triplets[key] = (max(existing_score, score), existing_relations, p_id)
+
+    # Build the text blocks
+    text_blocks = [
+        TextBlockMatch(
+            # XXX: Even though we are setting a paragraph_id, the text is not coming from the paragraph
+            paragraph_id=p_id,
+            scores=[GraphScore(score=score)],
+            score_type=SCORE_TYPE.RELATION_RELEVANCE,
+            order=0,
+            text=f"- {ent} {rel} {tail}",  # Manually build the text
+            position=TextPosition(
+                page_number=0,
+                index=0,
+                start=0,
+                end=0,
+                start_seconds=[],
+                end_seconds=[],
+            ),
+            field_labels=[],
+            paragraph_labels=[],
+            fuzzy_search=False,
+            is_a_table=False,
+            representation_file="",
+            page_with_visual=False,
+            relevant_relations=relations,
+        )
+        for (ent, rel, tail), (score, relations, p_id) in triplets.items()
+        if p_id is not None
+    ]
+    return text_blocks
+
+
+def get_paragraph_info_from_relations(
+    relations: Relations,
+    scores: dict[str, list[float]],
+) -> list[RelationsParagraphMatch]:
+    """
+    Gathers paragraph info from the 'relations' object, merges relations by paragraph,
+    and removes paragraphs contained entirely within others.
+    """
+
+    # Group paragraphs by field so we can detect containment
+    paragraphs_by_field: dict[FieldId, list[RelationsParagraphMatch]] = defaultdict(
+        list
+    )
+
+    # Loop over each entity in the relation graph
+    for ent, subgraph in relations.entities.items():
+        if ent not in scores:
+            continue
+        for rel_score, rel in zip(scores[ent], subgraph.related_to):
+            if rel.metadata and rel.metadata.paragraph_id:
+                p_id = ParagraphId.from_string(rel.metadata.paragraph_id)
+                match = RelationsParagraphMatch(
+                    paragraph_id=p_id,
+                    score=rel_score,
+                    relations=Relations(
+                        entities={ent: EntitySubgraph(related_to=[rel])}
+                    ),
+                )
+                paragraphs_by_field[p_id.field_id].append(match)
+
+    # For each field, sort paragraphs by start asc, end desc, and do one pass to remove contained ones
+    final_paragraphs: list[RelationsParagraphMatch] = []
+
+    for _, paragraph_list in paragraphs_by_field.items():
+        # Sort by paragraph_start ascending; if tie, paragraph_end descending
+        paragraph_list.sort(
+            key=lambda m: (
+                m.paragraph_id.paragraph_start,
+                -m.paragraph_id.paragraph_end,
+            )
+        )
+
+        kept: list[RelationsParagraphMatch] = []
+        current_max_end = -1
+
+        for match in paragraph_list:
+            end = match.paragraph_id.paragraph_end
+
+            # If end <= current_max_end, this paragraph is contained last one
+            if end <= current_max_end:
+                # We merge the scores and relations
+                container = kept[-1]
+                container.score = max(container.score, match.score)
+                for ent, subgraph in match.relations.entities.items():
+                    if ent not in container.relations.entities:
+                        container.relations.entities[ent] = EntitySubgraph(
+                            related_to=[]
+                        )
+                    container.relations.entities[ent].related_to.extend(
+                        subgraph.related_to
+                    )
+
+            else:
+                # Not contained; keep it
+                kept.append(match)
+                current_max_end = end
+        final_paragraphs.extend(kept)
+
+    return final_paragraphs
+
+
+async def build_graph_response(
+    *,
+    search_sdk: NucliaDBAsync,
+    kbid: str,
+    query: str,
+    final_relations: Relations,
+    scores: dict[str, list[float]],
+    top_k: int,
+    reranker: Reranker,
+    relation_text_as_paragraphs: bool,
+    show: list[ResourceProperties] = [],
+    extracted: list[ExtractedDataTypeName] = [],
+    field_type_filter: list[FieldTypeName] = [],
+) -> KnowledgeboxFindResults:
+    if relation_text_as_paragraphs:
+        text_blocks = build_text_blocks_from_relations(final_relations, scores)
+    else:
+        paragraphs_info = get_paragraph_info_from_relations(final_relations, scores)
+        text_blocks = relations_matches_to_text_block_matches(paragraphs_info)
+
+    # hydrate and rerank
+    resource_hydration_options = ResourceHydrationOptions(
+        show=show, extracted=extracted, field_type_filter=field_type_filter
+    )
+    text_block_hydration_options = TextBlockHydrationOptions(only_hydrate_empty=True)
+    reranking_options = RerankingOptions(kbid=kbid, query=query)
+    text_blocks, resources, best_matches = await hydrate_and_rerank(
+        search_sdk,
+        text_blocks,
+        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,
+    )
+
+    find_resources = compose_find_resources(text_blocks, resources)
+
+    return KnowledgeboxFindResults(
+        query=query,
+        resources=find_resources,
+        best_matches=best_matches,
+        relations=final_relations,
+        total=len(text_blocks),
+    )
+
+
+def relations_match_to_text_block_match(
+    paragraph_match: RelationsParagraphMatch,
+) -> TextBlockMatch:
+    """
+    Given a paragraph_id, return a TextBlockMatch with the bare minimum fields
+    This is required by the Graph Strategy to get text blocks from the relevant paragraphs
+    """
+    # XXX: this is a workaround for the fact we always assume retrieval means keyword/semantic search and
+    # the hydration and find response building code works with TextBlockMatch, we extended it to have relevant relations information
+    parsed_paragraph_id = paragraph_match.paragraph_id
+    return TextBlockMatch(
+        paragraph_id=parsed_paragraph_id,
+        scores=[GraphScore(score=paragraph_match.score)],
+        score_type=SCORE_TYPE.RELATION_RELEVANCE,
+        order=0,  # NOTE: this will be filled later
+        text="",  # NOTE: this will be filled later too
+        position=TextPosition(
+            page_number=0,
+            index=0,
+            start=parsed_paragraph_id.paragraph_start,
+            end=parsed_paragraph_id.paragraph_end,
+            start_seconds=[],
+            end_seconds=[],
+        ),
+        field_labels=[],
+        paragraph_labels=[],
+        fuzzy_search=False,
+        is_a_table=False,
+        representation_file="",
+        page_with_visual=False,
+        relevant_relations=paragraph_match.relations,
+    )
+
+
+def relations_matches_to_text_block_matches(
+    paragraph_matches: Collection[RelationsParagraphMatch],
+) -> list[TextBlockMatch]:
+    return [relations_match_to_text_block_match(match) for match in paragraph_matches]
+
+
+async def find_graph_neighbours(
+    search_sdk: NucliaDBAsync,
+    kbid: str,
+    entities_to_explore: list[RelationNode],
+    explored_entities: set[FrozenRelationNode],
+    exclude_processor_relations: bool,
+) -> GraphSearchResponse:
+    subqueries: list[GraphPathQuery] = []
+
+    # Explore starting from some entities
+    query_to_explore = [
+        GraphPath(
+            source=GraphNode(
+                value=entity.value,
+                type=RelationNodeTypePbMap[entity.ntype],
+                group=entity.subtype,
+            ),
+            destination=GraphNode(
+                type=RelationNodeType.ENTITY,
+            ),
+            undirected=True,
+        )
+        for entity in entities_to_explore
+    ]
+    if len(query_to_explore) > 0:
+        subqueries.append(Or(operands=query_to_explore))
+
+    # Do not return already known entities
+    if explored_entities:
+        query_exclude_explored = [
+            GraphPath(
+                source=GraphNode(
+                    value=explored.value,
+                    type=RelationNodeTypePbMap[explored.ntype],
+                    group=explored.subtype,
+                ),
+                undirected=True,
+            )
+            for explored in explored_entities
+        ]
+        if len(query_exclude_explored):
+            subqueries.append(Not(operand=Or(operands=query_exclude_explored)))
+
+    # Only include relations between entities
+    only_entities = Relation(type=RelationType.ENTITY)
+    subqueries.append(only_entities)
+
+    # Exclude processor entities
+    if exclude_processor_relations:
+        exclude_processor: GraphPathQuery = Not(
+            operand=Generated(by=Generator.PROCESSOR)
+        )
+        subqueries.append(exclude_processor)
+
+    graph_query = GraphSearchRequest(query=And(operands=subqueries), top_k=100)
+    response = await rpc.graph_paths(search_sdk, kbid, graph_query)
+
+    return response
+
+
+def merge_relations_results(
+    graph_response: GraphSearchResponse,
+    query_entry_points: Iterable[RelationNode],
+    only_with_metadata: bool = False,
+) -> Relations:
+    """Merge relation search responses into a single Relations object while applying filters."""
+    relations = Relations(entities={})
+
+    for entry_point in query_entry_points:
+        relations.entities[entry_point.value] = EntitySubgraph(related_to=[])
+
+    for path in graph_response.paths:
+        relation = path.relation
+        source = path.source
+        destination = path.destination
+        metadata = path.metadata
+
+        # BUG: if for some reason we don't have the field_id and we skip the
+        # next if, we'll reuse a wrong resource id
+        if path.metadata is not None and path.metadata.field_id:
+            resource_id = path.metadata.field_id.split("/")[0]
+
+        if only_with_metadata and path.metadata is None:
+            continue
+
+        if source.value in relations.entities:
+            relations.entities[source.value].related_to.append(
+                DirectionalRelation(
+                    entity=destination.value,
+                    entity_type=destination.type,
+                    entity_subtype=destination.group,
+                    relation=relation.type,
+                    relation_label=relation.label,
+                    direction=RelationDirection.OUT,
+                    metadata=RelationMetadata(
+                        paragraph_id=metadata.paragraph_id,
+                        # TODO(decoupled-ask): do we need more fields in the relation metadata?
+                    )
+                    if metadata is not None
+                    else None,
+                    resource_id=resource_id,
+                )
+            )
+
+        elif destination.value in relations.entities:
+            relations.entities[destination.value].related_to.append(
+                DirectionalRelation(
+                    entity=source.value,
+                    entity_type=source.type,
+                    entity_subtype=source.group,
+                    relation=relation.type,
+                    relation_label=relation.label,
+                    direction=RelationDirection.IN,
+                    metadata=RelationMetadata(
+                        paragraph_id=metadata.paragraph_id,
+                        # TODO(decoupled-ask): do we need more fields in the relation metadata?
+                    )
+                    if metadata is not None
+                    else None,
+                    resource_id=resource_id,
+                )
+            )
+
+    return relations