k3-node 1.0.0__py3-none-any.whl

k3_node/rag/subgraph.py

@@ -0,0 +1,270 @@
+"""Subgraph extraction utilities for GraphRAG and Knowledge Graph LLM integration."""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from keras import ops
+
+from k3_node.data import Data
+
+
+@dataclass
+class SubgraphResult:
+    """Structured result of a subgraph extraction around retrieved entities.
+
+    Attributes:
+        edge_index: Tensor of shape `(2, num_edges)` with subgraph edges.
+        edge_type: Optional tensor of shape `(num_edges,)` with relation types.
+        edge_attr: Optional tensor of shape `(num_edges, edge_dim)` with edge features.
+        x: Optional tensor of shape `(num_nodes, in_channels)` with node features.
+        nodes: 1D numpy array of original node indices in the full graph.
+        center_nodes: 1D numpy array of center entity indices in the relabeled subgraph.
+        mapping: Dictionary mapping original node index to subgraph node index.
+        edge_mask: 1D boolean array indicating edges kept from original graph.
+        num_nodes: Number of nodes in the extracted subgraph.
+        num_edges: Number of edges in the extracted subgraph.
+    """
+
+    edge_index: any
+    edge_type: Optional[any] = None
+    edge_attr: Optional[any] = None
+    x: Optional[any] = None
+    nodes: Optional[np.ndarray] = None
+    center_nodes: Optional[np.ndarray] = None
+    mapping: Optional[Dict[int, int]] = None
+    edge_mask: Optional[np.ndarray] = None
+    num_nodes: int = 0
+    num_edges: int = 0
+
+    def to_data(self) -> Data:
+        """Convert extracted subgraph into a K3-Node Data object."""
+        return Data(
+            x=self.x,
+            edge_index=self.edge_index,
+            edge_type=self.edge_type,
+            edge_attr=self.edge_attr,
+            center_nodes=ops.convert_to_tensor(self.center_nodes, dtype="int32")
+            if self.center_nodes is not None
+            else None,
+            original_nodes=ops.convert_to_tensor(self.nodes, dtype="int32")
+            if self.nodes is not None
+            else None,
+        )
+
+
+def extract_subgraph(
+    entities: Union[int, Sequence[int], np.ndarray, any],
+    edge_index: Union[np.ndarray, any],
+    edge_type: Optional[Union[np.ndarray, any]] = None,
+    edge_attr: Optional[Union[np.ndarray, any]] = None,
+    x: Optional[Union[np.ndarray, any]] = None,
+    num_hops: int = 2,
+    max_nodes_per_hop: Optional[int] = None,
+    directed: bool = False,
+    relabel_nodes: bool = True,
+    num_nodes: Optional[int] = None,
+) -> SubgraphResult:
+    """Extract multi-hop enclosing or ego-subgraphs around retrieved entities.
+
+    Given a knowledge graph or relational graph, this function expands `num_hops`
+    around seed `entities`, keeping all induced edges, relation types, and node/edge
+    attributes.
+
+    Args:
+        entities: Single node index or list/array of seed entity indices.
+        edge_index: Graph connectivity tensor of shape `(2, num_edges)`.
+        edge_type: Optional 1D relation type tensor of shape `(num_edges,)`.
+        edge_attr: Optional edge attribute tensor of shape `(num_edges, edge_dim)`.
+        x: Optional node feature tensor of shape `(total_nodes, in_channels)`.
+        num_hops: Number of hops to expand around retrieved entities. (default: 2)
+        max_nodes_per_hop: Optional maximum number of neighboring nodes to keep
+            per hop (useful to restrict explosion on hub entities).
+        directed: If True, only follows outgoing edges. If False, follows edges
+            in both directions (standard for KG context expansion). (default: False)
+        relabel_nodes: If True, relabels subgraph node IDs to `0..num_subgraph_nodes-1`.
+            (default: True)
+        num_nodes: Optional total number of nodes in graph. Inferred if not given.
+
+    Returns:
+        `SubgraphResult` containing relabeled edge_index, edge_type, features,
+        and entity mappings.
+    """
+    edge_index_np = np.asarray(ops.convert_to_numpy(edge_index)).astype(np.int64)
+    if num_nodes is None:
+        num_nodes = int(edge_index_np.max()) + 1 if edge_index_np.size > 0 else 0
+
+    entities_arr = np.atleast_1d(np.asarray(ops.convert_to_numpy(entities))).astype(np.int64)
+    if entities_arr.size == 0:
+        empty_ei = ops.convert_to_tensor(np.zeros((2, 0), dtype=np.int64), dtype=edge_index.dtype)
+        return SubgraphResult(
+            edge_index=empty_ei,
+            nodes=np.array([], dtype=np.int64),
+            center_nodes=np.array([], dtype=np.int64),
+            mapping={},
+            num_nodes=0,
+            num_edges=0,
+        )
+
+    row, col = edge_index_np[0], edge_index_np[1]
+    subsets = [entities_arr]
+    visited = set(entities_arr.tolist())
+
+    for _ in range(num_hops):
+        current_frontier = subsets[-1]
+        if len(current_frontier) == 0:
+            break
+
+        mask_frontier = np.zeros(num_nodes, dtype=bool)
+        mask_frontier[current_frontier] = True
+
+        # If undirected (or GraphRAG bidirectional expansion), collect neighbors in both directions
+        if not directed:
+            neighbors_out = col[mask_frontier[row]]
+            neighbors_in = row[mask_frontier[col]]
+            new_neighbors = np.concatenate([neighbors_out, neighbors_in])
+        else:
+            new_neighbors = col[mask_frontier[row]]
+
+        if new_neighbors.size > 0:
+            unique_new = np.unique(new_neighbors)
+            unseen = [n for n in unique_new if n not in visited]
+            if max_nodes_per_hop is not None and len(unseen) > max_nodes_per_hop:
+                unseen = unseen[:max_nodes_per_hop]
+            visited.update(unseen)
+            subsets.append(np.array(unseen, dtype=np.int64))
+        else:
+            break
+
+    subgraph_nodes = np.unique(np.concatenate([s for s in subsets if len(s) > 0]))
+
+    # Induced edge mask: both endpoints must be in subgraph_nodes
+    node_mask = np.zeros(num_nodes, dtype=bool)
+    node_mask[subgraph_nodes] = True
+    edge_mask = node_mask[row] & node_mask[col]
+
+    sub_edge_index = edge_index_np[:, edge_mask]
+
+    mapping = {int(orig): int(new_idx) for new_idx, orig in enumerate(subgraph_nodes)}
+    center_indices = np.array([mapping[int(e)] for e in entities_arr if int(e) in mapping], dtype=np.int64)
+
+    if relabel_nodes:
+        new_id_table = np.full(num_nodes, -1, dtype=np.int64)
+        new_id_table[subgraph_nodes] = np.arange(len(subgraph_nodes), dtype=np.int64)
+        sub_edge_index = new_id_table[sub_edge_index]
+
+    # Preserve tensors in their original backend / dtype
+    sub_edge_index_tensor = ops.convert_to_tensor(sub_edge_index, dtype=edge_index.dtype)
+
+    sub_edge_type_tensor = None
+    if edge_type is not None:
+        et_np = np.asarray(ops.convert_to_numpy(edge_type))[edge_mask]
+        sub_edge_type_tensor = ops.convert_to_tensor(et_np, dtype=edge_type.dtype)
+
+    sub_edge_attr_tensor = None
+    if edge_attr is not None:
+        ea_np = np.asarray(ops.convert_to_numpy(edge_attr))[edge_mask]
+        sub_edge_attr_tensor = ops.convert_to_tensor(ea_np, dtype=edge_attr.dtype)
+
+    sub_x_tensor = None
+    if x is not None:
+        x_np = np.asarray(ops.convert_to_numpy(x))[subgraph_nodes]
+        sub_x_tensor = ops.convert_to_tensor(x_np, dtype=x.dtype)
+
+    return SubgraphResult(
+        edge_index=sub_edge_index_tensor,
+        edge_type=sub_edge_type_tensor,
+        edge_attr=sub_edge_attr_tensor,
+        x=sub_x_tensor,
+        nodes=subgraph_nodes,
+        center_nodes=center_indices,
+        mapping=mapping,
+        edge_mask=edge_mask,
+        num_nodes=len(subgraph_nodes),
+        num_edges=int(sub_edge_index.shape[1]),
+    )
+
+
+class KGEntityRetriever:
+    """Knowledge Graph Entity and Subgraph Retriever for GraphRAG.
+
+    Maintains entity name dictionaries and relation mappings, extracts seed entities
+    from text queries, and retrieves enclosing multi-hop subgraphs.
+
+    Args:
+        entity_to_id: Dictionary mapping entity strings to node integer IDs.
+        relation_to_id: Dictionary mapping relation strings to edge_type integer IDs.
+        edge_index: Graph connectivity tensor of shape `(2, num_edges)`.
+        edge_type: Optional relation type tensor of shape `(num_edges,)`.
+        edge_attr: Optional edge feature tensor.
+        x: Optional node feature tensor.
+    """
+
+    def __init__(
+        self,
+        entity_to_id: Dict[str, int],
+        relation_to_id: Dict[str, int],
+        edge_index: Union[np.ndarray, any],
+        edge_type: Optional[Union[np.ndarray, any]] = None,
+        edge_attr: Optional[Union[np.ndarray, any]] = None,
+        x: Optional[Union[np.ndarray, any]] = None,
+    ):
+        self.entity_to_id = entity_to_id
+        self.relation_to_id = relation_to_id
+        self.id_to_entity = {v: k for k, v in entity_to_id.items()}
+        self.id_to_relation = {v: k for k, v in relation_to_id.items()}
+
+        self.edge_index = edge_index
+        self.edge_type = edge_type
+        self.edge_attr = edge_attr
+        self.x = x
+
+    def get_entity_id(self, name: str) -> Optional[int]:
+        """Look up entity ID by exact name (case-insensitive fallback)."""
+        if name in self.entity_to_id:
+            return self.entity_to_id[name]
+        # Case-insensitive fallback
+        lower_map = {k.lower(): v for k, v in self.entity_to_id.items()}
+        return lower_map.get(name.lower(), None)
+
+    def find_entities_in_text(self, text: str) -> List[str]:
+        """Find matching known entity names mentioned in a query text."""
+        text_lower = text.lower()
+        matched = []
+        # Sort by length descending to match longest phrases first
+        for name in sorted(self.entity_to_id.keys(), key=lambda s: len(s), reverse=True):
+            if name.lower() in text_lower:
+                matched.append(name)
+        return matched
+
+    def retrieve_subgraph(
+        self,
+        entities: Union[Sequence[Union[str, int]], str, int],
+        num_hops: int = 2,
+        max_nodes_per_hop: Optional[int] = None,
+        directed: bool = False,
+    ) -> SubgraphResult:
+        """Extract multi-hop subgraph around the specified entity names or IDs."""
+        if isinstance(entities, (str, int)):
+            entities = [entities]
+
+        entity_ids = []
+        for e in entities:
+            if isinstance(e, str):
+                eid = self.get_entity_id(e)
+                if eid is not None:
+                    entity_ids.append(eid)
+            else:
+                entity_ids.append(int(e))
+
+        return extract_subgraph(
+            entities=entity_ids,
+            edge_index=self.edge_index,
+            edge_type=self.edge_type,
+            edge_attr=self.edge_attr,
+            x=self.x,
+            num_hops=num_hops,
+            max_nodes_per_hop=max_nodes_per_hop,
+            directed=directed,
+            relabel_nodes=True,
+        )

k3_node/rag/test_rag.py

@@ -0,0 +1,347 @@
+"""Tests for k3_node.rag GraphRAG and KG-LLM connectors."""
+
+import numpy as np
+import pytest
+from keras import ops
+
+import k3_node as k3
+from k3_node.layers.kge import TransE
+from k3_node.rag import (
+    GraphPrefixProjector,
+    GraphRAG,
+    KGEntityRetriever,
+    KGLLMConnector,
+    RGCNSubGraphEncoder,
+    SubgraphResult,
+    TransEPrefixEncoder,
+    extract_subgraph,
+    format_llm_prompt,
+    subgraph_to_triples,
+    verbalize_subgraph,
+)
+
+
+@pytest.fixture
+def sample_kg():
+    """Create a sample knowledge graph for testing.
+
+    Graph:
+        0 (Aspirin) --[0: treats]--> 1 (Headache)
+        0 (Aspirin) --[1: inhibits]--> 2 (COX-1)
+        2 (COX-1) --[2: produces]--> 3 (Prostaglandin)
+        3 (Prostaglandin) --[3: causes]--> 1 (Headache)
+        4 (Ibuprofen) --[0: treats]--> 1 (Headache)
+        4 (Ibuprofen) --[1: inhibits]--> 2 (COX-1)
+    """
+    edges = np.array(
+        [
+            [0, 0, 2, 3, 4, 4],
+            [1, 2, 3, 1, 1, 2],
+        ],
+        dtype=np.int64,
+    )
+    edge_type = np.array([0, 1, 2, 3, 0, 1], dtype=np.int64)
+    x = np.random.randn(5, 16).astype(np.float32)
+
+    entity_to_id = {
+        "Aspirin": 0,
+        "Headache": 1,
+        "COX-1": 2,
+        "Prostaglandin": 3,
+        "Ibuprofen": 4,
+    }
+    relation_to_id = {
+        "treats": 0,
+        "inhibits": 1,
+        "produces": 2,
+        "causes": 3,
+    }
+
+    return {
+        "edge_index": ops.convert_to_tensor(edges, dtype="int32"),
+        "edge_type": ops.convert_to_tensor(edge_type, dtype="int32"),
+        "x": ops.convert_to_tensor(x, dtype="float32"),
+        "entity_to_id": entity_to_id,
+        "relation_to_id": relation_to_id,
+        "num_nodes": 5,
+        "num_relations": 4,
+    }
+
+
+def test_extract_subgraph_1hop(sample_kg):
+    sub = extract_subgraph(
+        entities=[0],  # Aspirin
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        x=sample_kg["x"],
+        num_hops=1,
+        directed=False,
+    )
+
+    assert isinstance(sub, SubgraphResult)
+    # Neighbors of 0 within 1 hop: 0, 1, 2
+    assert set(sub.nodes.tolist()) == {0, 1, 2}
+    assert sub.center_nodes.tolist() == [sub.mapping[0]]
+    assert sub.num_nodes == 3
+    assert sub.num_edges > 0
+    assert sub.x is not None
+    assert ops.shape(sub.x)[0] == 3
+
+    # Check to_data()
+    data = sub.to_data()
+    assert isinstance(data, k3.data.Data)
+    assert ops.shape(data.edge_index)[0] == 2
+
+
+def test_extract_subgraph_2hop_multiseed(sample_kg):
+    sub = extract_subgraph(
+        entities=[0, 3],  # Aspirin & Prostaglandin
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        num_hops=2,
+    )
+    assert sub.num_nodes == 5
+    assert len(sub.center_nodes) == 2
+
+
+def test_kg_entity_retriever(sample_kg):
+    retriever = KGEntityRetriever(
+        entity_to_id=sample_kg["entity_to_id"],
+        relation_to_id=sample_kg["relation_to_id"],
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        x=sample_kg["x"],
+    )
+
+    # Lookup
+    assert retriever.get_entity_id("Aspirin") == 0
+    assert retriever.get_entity_id("aspirin") == 0  # case-insensitive
+    assert retriever.get_entity_id("Unknown") is None
+
+    # Text entity extraction
+    query = "Does Aspirin or Ibuprofen treat headache?"
+    found = retriever.find_entities_in_text(query)
+    assert "Aspirin" in found
+    assert "Ibuprofen" in found
+    assert "Headache" in found
+
+    # Subgraph retrieval
+    sub = retriever.retrieve_subgraph(["Aspirin"], num_hops=1)
+    assert sub.num_nodes >= 2
+
+
+def test_verbalization(sample_kg):
+    sub = extract_subgraph(
+        entities=[0],
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        num_hops=1,
+    )
+
+    id_to_e = {v: k for k, v in sample_kg["entity_to_id"].items()}
+    id_to_r = {v: k for k, v in sample_kg["relation_to_id"].items()}
+
+    # Triples list
+    triples = subgraph_to_triples(sub, id_to_e, id_to_r)
+    assert len(triples) > 0
+    assert ("Aspirin", "treats", "Headache") in triples
+
+    # Markdown format
+    md_text = verbalize_subgraph(sub, id_to_e, id_to_r, format_style="markdown")
+    assert "**Aspirin**" in md_text
+    assert "*treats*" in md_text
+
+    # Natural language format
+    natural_text = verbalize_subgraph(sub, id_to_e, id_to_r, format_style="natural")
+    assert "Aspirin treats Headache." in natural_text
+
+
+def test_prompt_formatting():
+    context = "- **Aspirin** — *treats* -> **Headache**"
+    query = "What treats headache?"
+
+    # Llama 3
+    llama3_prompt = format_llm_prompt(query, context, model_family="llama3")
+    assert "<|start_header_id|>system<|end_header_id|>" in llama3_prompt
+    assert "<|start_header_id|>user<|end_header_id|>" in llama3_prompt
+    assert "Knowledge Graph Context:" in llama3_prompt
+    assert query in llama3_prompt
+
+    # Mistral
+    mistral_prompt = format_llm_prompt(query, context, model_family="mistral")
+    assert "<s>[INST]" in mistral_prompt
+    assert "[/INST]" in mistral_prompt
+
+    # ChatML
+    chatml_prompt = format_llm_prompt(query, context, model_family="chatml")
+    assert "<|im_start|>system" in chatml_prompt
+
+
+def test_rgcn_subgraph_encoder(sample_kg):
+    encoder = RGCNSubGraphEncoder(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=64,
+        num_relations=sample_kg["num_relations"],
+        num_layers=2,
+        pooling="center",
+    )
+
+    sub = extract_subgraph(
+        entities=[0],
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        x=sample_kg["x"],
+        num_hops=1,
+    )
+
+    emb = encoder.encode_subgraph(sub)
+    assert ops.shape(emb) == (1, 64)
+
+    # Test other pooling options
+    encoder_mean = RGCNSubGraphEncoder(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=64,
+        num_relations=sample_kg["num_relations"],
+        pooling="mean",
+    )
+    emb_mean = encoder_mean.encode_subgraph(sub)
+    assert ops.shape(emb_mean) == (1, 64)
+
+
+def test_transe_prefix_encoder(sample_kg):
+    transe = TransE(
+        num_nodes=sample_kg["num_nodes"],
+        num_relations=sample_kg["num_relations"],
+        hidden_channels=32,
+    )
+
+    kge_encoder = TransEPrefixEncoder(
+        kge_model=transe,
+        out_channels=64,
+        pooling="center",
+    )
+
+    sub = extract_subgraph(
+        entities=[0, 1],
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        num_hops=1,
+    )
+
+    emb = kge_encoder.encode_subgraph(sub)
+    assert ops.shape(emb) == (1, 64)
+
+    # Test standalone initialization without pretrained model
+    standalone_encoder = TransEPrefixEncoder(
+        num_nodes=10,
+        num_relations=4,
+        embedding_dim=32,
+        out_channels=64,
+    )
+    emb_standalone = standalone_encoder.encode_subgraph(sub)
+    assert ops.shape(emb_standalone) == (1, 64)
+
+
+def test_graph_prefix_projector():
+    projector = GraphPrefixProjector(
+        in_channels=64,
+        llm_dim=256,  # test dimension
+        num_prefix_tokens=4,
+        projector_type="mlp",
+    )
+
+    graph_emb = ops.convert_to_tensor(np.random.randn(1, 64).astype(np.float32))
+    prefix = projector(graph_emb)
+    assert ops.shape(prefix) == (1, 4, 256)
+
+    # Linear projector
+    lin_projector = GraphPrefixProjector(
+        in_channels=64,
+        llm_dim=256,
+        num_prefix_tokens=4,
+        projector_type="linear",
+    )
+    prefix_lin = lin_projector(graph_emb)
+    assert ops.shape(prefix_lin) == (1, 4, 256)
+
+
+def test_kg_llm_connector(sample_kg):
+    encoder = RGCNSubGraphEncoder(
+        in_channels=16,
+        hidden_channels=32,
+        out_channels=64,
+        num_relations=sample_kg["num_relations"],
+    )
+    projector = GraphPrefixProjector(
+        in_channels=64,
+        llm_dim=512,
+        num_prefix_tokens=4,
+    )
+    connector = KGLLMConnector(encoder=encoder, projector=projector)
+
+    sub = extract_subgraph(
+        entities=[0],
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        x=sample_kg["x"],
+        num_hops=1,
+    )
+
+    prefix_tokens = connector.encode_subgraph(sub)
+    assert ops.shape(prefix_tokens) == (1, 4, 512)
+
+    # Test prefix injection into LLM text tokens
+    text_tokens = ops.convert_to_tensor(np.random.randn(1, 10, 512).astype(np.float32))
+    augmented = connector.inject_prefix(text_tokens, prefix_tokens)
+    assert ops.shape(augmented) == (1, 14, 512)
+
+    # Test attention mask extension
+    attn_mask = ops.ones((1, 10), dtype="int32")
+    extended_mask = connector.extend_attention_mask(attn_mask, num_prefix_tokens=4)
+    assert ops.shape(extended_mask) == (1, 14)
+
+
+def test_graph_rag_pipeline(sample_kg):
+    # Pipeline with RGCN
+    rag_rgcn = GraphRAG(
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        entity_to_id=sample_kg["entity_to_id"],
+        relation_to_id=sample_kg["relation_to_id"],
+        x=sample_kg["x"],
+        encoder_type="rgcn",
+        hidden_dim=32,
+        encoder_out_dim=64,
+        llm_dim=256,
+        num_prefix_tokens=4,
+    )
+
+    # Retrieval
+    sub = rag_rgcn.retrieve(["Aspirin"], num_hops=1)
+    assert sub.num_nodes >= 2
+
+    # Prompt builder
+    prompt = rag_rgcn.build_prompt("What does Aspirin treat?", model_family="llama3")
+    assert "Aspirin" in prompt
+    assert "<|begin_of_text|>" in prompt
+
+    # Prefix encoding
+    prefix = rag_rgcn.encode_prefix(sub)
+    assert ops.shape(prefix) == (1, 4, 256)
+
+    # Pipeline with TransE
+    rag_transe = GraphRAG(
+        edge_index=sample_kg["edge_index"],
+        edge_type=sample_kg["edge_type"],
+        entity_to_id=sample_kg["entity_to_id"],
+        relation_to_id=sample_kg["relation_to_id"],
+        encoder_type="transe",
+        hidden_dim=32,
+        encoder_out_dim=64,
+        llm_dim=256,
+        num_prefix_tokens=4,
+    )
+    prefix_transe = rag_transe.encode_prefix(sub)
+    assert ops.shape(prefix_transe) == (1, 4, 256)