k3-node 1.0.0__py3-none-any.whl

k3_node/rag/encoders.py

@@ -0,0 +1,312 @@
+"""GNN and KGE Subgraph Encoders for GraphRAG and KG-LLM augmentation."""
+
+from typing import List, Literal, Optional, Tuple, Union
+
+import keras
+from keras import layers, ops
+
+from k3_node.layers.conv import GCNConv, RGCNConv, SAGEConv
+from k3_node.layers.kge import KGEModel, TransE
+from k3_node.rag.subgraph import SubgraphResult
+
+
+class RGCNSubGraphEncoder(keras.Model):
+    r"""Relational Graph Convolutional Network (RGCN) encoder for multi-relational subgraphs.
+
+    Processes subgraphs extracted from Knowledge Graphs with multiple relation types,
+    updating entity representations through relational message passing and pooling them
+    into a dense graph embedding.
+
+    Args:
+        in_channels: Dimensionality of input node features.
+        hidden_channels: Hidden representation dimension.
+        out_channels: Output graph embedding dimension.
+        num_relations: Total number of relation types in the Knowledge Graph.
+        num_layers: Number of RGCN message passing layers. (default: 2)
+        num_bases: Optional number of basis decomposition components for relation weights.
+        pooling: Readout pooling strategy across subgraph nodes:
+            - `"mean"`: Global average over all subgraph nodes.
+            - `"sum"`: Global sum over all subgraph nodes.
+            - `"max"`: Global maximum over all subgraph nodes.
+            - `"center"`: Pool only the retrieved center seed entities.
+            - `"none"`: Return all node embeddings without pooling.
+        dropout: Dropout rate applied between convolution layers. (default: 0.0)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_relations: int,
+        num_layers: int = 2,
+        num_bases: Optional[int] = None,
+        pooling: Literal["mean", "sum", "max", "center", "none"] = "mean",
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_relations = num_relations
+        self.num_layers = num_layers
+        self.pooling = pooling
+        self.dropout_rate = dropout
+
+        self.convs = []
+        for i in range(num_layers):
+            c_in = in_channels if i == 0 else hidden_channels
+            c_out = out_channels if i == num_layers - 1 else hidden_channels
+            self.convs.append(
+                RGCNConv(
+                    in_channels=c_in,
+                    out_channels=c_out,
+                    num_relations=num_relations,
+                    num_bases=num_bases,
+                )
+            )
+
+        self.act = layers.Activation("relu")
+        self.drop = layers.Dropout(dropout) if dropout > 0.0 else None
+
+    def call(
+        self,
+        x: any,
+        edge_index: any,
+        edge_type: Optional[any] = None,
+        center_nodes: Optional[any] = None,
+        training: Optional[bool] = None,
+    ) -> any:
+        """Forward pass encoding the subgraph into a graph embedding.
+
+        Args:
+            x: Node feature tensor of shape `(num_nodes, in_channels)`.
+            edge_index: Graph edge indices of shape `(2, num_edges)`.
+            edge_type: 1D tensor of relation IDs for each edge `(num_edges,)`.
+            center_nodes: Optional 1D tensor of seed entity indices in the subgraph.
+            training: Whether running in training mode.
+
+        Returns:
+            Tensor of shape `(1, out_channels)` (or `(num_nodes, out_channels)` if `pooling="none"`).
+        """
+        h = x
+        for i, conv in enumerate(self.convs):
+            h = conv(h, edge_index, edge_type=edge_type)
+            if i < self.num_layers - 1:
+                h = self.act(h)
+                if self.drop is not None:
+                    h = self.drop(h, training=training)
+
+        if self.pooling == "none":
+            return h
+
+        if self.pooling == "center" and center_nodes is not None:
+            center_nodes = ops.convert_to_tensor(center_nodes, dtype="int32")
+            if ops.shape(center_nodes)[0] > 0:
+                center_h = ops.take(h, center_nodes, axis=0)
+                pooled = ops.mean(center_h, axis=0, keepdims=True)
+                return pooled
+
+        if self.pooling == "sum":
+            return ops.sum(h, axis=0, keepdims=True)
+        elif self.pooling == "max":
+            return ops.max(h, axis=0, keepdims=True)
+        else:  # "mean" or fallback
+            return ops.mean(h, axis=0, keepdims=True)
+
+    def encode_subgraph(self, subgraph: SubgraphResult, default_x_dim: Optional[int] = None) -> any:
+        """Helper to encode a SubgraphResult directly."""
+        x = subgraph.x
+        if x is None:
+            dim = default_x_dim if default_x_dim is not None else self.in_channels
+            num_n = subgraph.num_nodes if subgraph.num_nodes > 0 else 1
+            x = ops.ones((num_n, dim), dtype="float32")
+
+        return self(
+            x=x,
+            edge_index=subgraph.edge_index,
+            edge_type=subgraph.edge_type,
+            center_nodes=subgraph.center_nodes,
+        )
+
+
+class TransEPrefixEncoder(keras.Model):
+    r"""Knowledge Graph Embedding (KGE) prefix encoder using TransE representations.
+
+    Uses pretrained or end-to-end entity and relation embeddings from TransE ($h + r \approx t$)
+    to encode extracted knowledge subgraphs into a unified dense embedding vector.
+
+    Args:
+        num_nodes: Total number of entities in the knowledge graph.
+        num_relations: Total number of relation types.
+        embedding_dim: Dimension of entity and relation embeddings. (default: 64)
+        out_channels: Output projected graph embedding dimension. (default: 128)
+        kge_model: Optional pretrained `KGEModel` (e.g. `TransE`). If provided,
+            its embeddings are reused.
+        pooling: Readout pooling strategy across entities (`"mean"`, `"center"`, `"sum"`).
+    """
+
+    def __init__(
+        self,
+        num_nodes: Optional[int] = None,
+        num_relations: Optional[int] = None,
+        embedding_dim: int = 64,
+        out_channels: int = 128,
+        kge_model: Optional[KGEModel] = None,
+        pooling: Literal["mean", "sum", "center"] = "mean",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.embedding_dim = embedding_dim
+        self.out_channels = out_channels
+        self.pooling = pooling
+
+        if kge_model is not None:
+            self.node_emb = kge_model.node_emb
+            self.rel_emb = kge_model.rel_emb
+            self.embedding_dim = kge_model.hidden_channels
+        else:
+            if num_nodes is None or num_relations is None:
+                raise ValueError("Must provide either `kge_model` or both `num_nodes` and `num_relations`.")
+            self.node_emb = layers.Embedding(num_nodes, embedding_dim)
+            self.rel_emb = layers.Embedding(num_relations, embedding_dim)
+            self.node_emb.build((None,))
+            self.rel_emb.build((None,))
+
+        # Projection head combining entity representation and relational context
+        self.proj = keras.Sequential(
+            [
+                layers.Dense(out_channels, activation="relu"),
+                layers.LayerNormalization(),
+                layers.Dense(out_channels),
+            ]
+        )
+
+    def call(
+        self,
+        subgraph_nodes: any,
+        edge_index: Optional[any] = None,
+        edge_type: Optional[any] = None,
+        center_nodes: Optional[any] = None,
+    ) -> any:
+        """Encode a set of subgraph entities and relations into a dense embedding.
+
+        Args:
+            subgraph_nodes: 1D tensor of original entity IDs in the subgraph `(num_nodes,)`.
+            edge_index: Optional edge index tensor `(2, num_edges)`.
+            edge_type: Optional relation type tensor `(num_edges,)`.
+            center_nodes: Optional indices of seed entities within `subgraph_nodes`.
+
+        Returns:
+            Dense embedding tensor of shape `(1, out_channels)`.
+        """
+        subgraph_nodes = ops.convert_to_tensor(subgraph_nodes, dtype="int32")
+        entity_embeddings = self.node_emb(subgraph_nodes)  # [num_nodes, embedding_dim]
+
+        # Entity pooling
+        if self.pooling == "center" and center_nodes is not None:
+            center_idx = ops.convert_to_tensor(center_nodes, dtype="int32")
+            if ops.shape(center_idx)[0] > 0:
+                ent_h = ops.take(entity_embeddings, center_idx, axis=0)
+                node_repr = ops.mean(ent_h, axis=0, keepdims=True)
+            else:
+                node_repr = ops.mean(entity_embeddings, axis=0, keepdims=True)
+        elif self.pooling == "sum":
+            node_repr = ops.sum(entity_embeddings, axis=0, keepdims=True)
+        else:
+            node_repr = ops.mean(entity_embeddings, axis=0, keepdims=True)
+
+        # Relational translation context (h + r - t in TransE)
+        if edge_index is not None and edge_type is not None:
+            edge_index = ops.convert_to_tensor(edge_index, dtype="int32")
+            edge_type = ops.convert_to_tensor(edge_type, dtype="int32")
+            num_e = ops.shape(edge_index)[1]
+
+            if num_e > 0:
+                row, col = edge_index[0], edge_index[1]
+                h_sub = ops.take(entity_embeddings, row, axis=0)
+                t_sub = ops.take(entity_embeddings, col, axis=0)
+                r_sub = self.rel_emb(edge_type)
+                # TransE relation representation
+                triple_context = ops.mean(h_sub + r_sub - t_sub, axis=0, keepdims=True)
+                combined = ops.concatenate([node_repr, triple_context], axis=-1)
+            else:
+                combined = ops.concatenate([node_repr, ops.zeros_like(node_repr)], axis=-1)
+        else:
+            combined = ops.concatenate([node_repr, ops.zeros_like(node_repr)], axis=-1)
+
+        return self.proj(combined)
+
+    def encode_subgraph(self, subgraph: SubgraphResult) -> any:
+        """Helper to encode a SubgraphResult directly."""
+        nodes = subgraph.nodes if subgraph.nodes is not None else np.arange(subgraph.num_nodes)
+        return self(
+            subgraph_nodes=nodes,
+            edge_index=subgraph.edge_index,
+            edge_type=subgraph.edge_type,
+            center_nodes=subgraph.center_nodes,
+        )
+
+
+class GNNSubGraphEncoder(keras.Model):
+    r"""General-purpose GNN encoder (GCN / GraphSAGE) for homogeneous subgraphs.
+
+    Args:
+        in_channels: Input node feature dimensionality.
+        hidden_channels: Hidden representation dimension.
+        out_channels: Output graph embedding dimension.
+        conv_type: Type of GNN convolution (`"gcn"` or `"sage"`). (default: `"gcn"`)
+        num_layers: Number of convolution layers. (default: 2)
+        pooling: Pooling strategy (`"mean"`, `"sum"`, `"max"`, `"center"`).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        conv_type: Literal["gcn", "sage"] = "gcn",
+        num_layers: int = 2,
+        pooling: Literal["mean", "sum", "max", "center"] = "mean",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.pooling = pooling
+        self.num_layers = num_layers
+
+        ConvCls = GCNConv if conv_type.lower() == "gcn" else SAGEConv
+        self.convs = []
+        for i in range(num_layers):
+            c_in = in_channels if i == 0 else hidden_channels
+            c_out = out_channels if i == num_layers - 1 else hidden_channels
+            self.convs.append(ConvCls(c_in, c_out))
+
+        self.act = layers.Activation("relu")
+
+    def call(
+        self,
+        x: any,
+        edge_index: any,
+        center_nodes: Optional[any] = None,
+    ) -> any:
+        h = x
+        for i, conv in enumerate(self.convs):
+            h = conv(h, edge_index)
+            if i < self.num_layers - 1:
+                h = self.act(h)
+
+        if self.pooling == "center" and center_nodes is not None:
+            center_nodes = ops.convert_to_tensor(center_nodes, dtype="int32")
+            if ops.shape(center_nodes)[0] > 0:
+                return ops.mean(ops.take(h, center_nodes, axis=0), axis=0, keepdims=True)
+
+        if self.pooling == "sum":
+            return ops.sum(h, axis=0, keepdims=True)
+        elif self.pooling == "max":
+            return ops.max(h, axis=0, keepdims=True)
+        else:
+            return ops.mean(h, axis=0, keepdims=True)

k3_node/rag/pipeline.py

@@ -0,0 +1,192 @@
+"""High-level GraphRAG Pipeline integrating retrieval, verbalization, and LLM prefix injection."""
+
+from typing import Dict, List, Literal, Optional, Sequence, Union
+import numpy as np
+from keras import ops
+
+from k3_node.rag.encoders import RGCNSubGraphEncoder, TransEPrefixEncoder
+from k3_node.rag.projector import GraphPrefixProjector, KGLLMConnector
+from k3_node.rag.subgraph import KGEntityRetriever, SubgraphResult, extract_subgraph
+from k3_node.rag.verbalizer import format_llm_prompt, verbalize_subgraph
+
+
+class GraphRAG:
+    r"""End-to-End Graph-Augmented Generation (GraphRAG) Pipeline for Knowledge Graphs.
+
+    Provides a unified API for:
+    1. Extracting subgraphs around query entities.
+    2. Verbalizing structured facts into markdown or natural language prompts for Llama 3 / Mistral.
+    3. Projecting GNN (RGCN) or KGE (TransE) embeddings into dense prefix vectors for soft prompt augmentation.
+
+    Example:
+        ```python
+        import k3_node as k3
+        from k3_node.rag import GraphRAG
+
+        # Setup GraphRAG with your knowledge graph
+        rag = GraphRAG(
+            edge_index=edge_index,
+            edge_type=edge_type,
+            entity_to_id={"Aspirin": 0, "Headache": 1, "COX-1": 2},
+            relation_to_id={"treats": 0, "inhibits": 1},
+            llm_dim=4096,  # Llama 3 / Mistral embedding dimension
+            num_prefix_tokens=4,
+        )
+
+        # 1. Text-based GraphRAG prompt generation
+        prompt = rag.build_prompt(
+            query="How does Aspirin alleviate headache?",
+            entities=["Aspirin", "Headache"],
+            model_family="llama3",
+        )
+
+        # 2. Dense Prefix Vector encoding
+        subgraph = rag.retrieve(["Aspirin"], num_hops=2)
+        prefix_embeds = rag.encode_prefix(subgraph)  # [1, 4, 4096]
+        ```
+
+    Args:
+        edge_index: Graph connectivity tensor of shape `(2, num_edges)`.
+        edge_type: Optional relation IDs tensor of shape `(num_edges,)`.
+        entity_to_id: Dictionary mapping entity strings to node IDs.
+        relation_to_id: Optional dictionary mapping relation strings to relation IDs.
+        x: Optional node feature tensor.
+        num_relations: Optional total number of relation types.
+        encoder_type: Subgraph encoder architecture (`"rgcn"`, `"transe"`, or a custom encoder).
+        hidden_dim: Hidden dimension for GNN/KGE encoder. (default: 64)
+        encoder_out_dim: Output dimension of graph encoder before LLM projection. (default: 128)
+        llm_dim: Target LLM embedding dimension (e.g. 4096 for Llama 3 / Mistral). (default: 4096)
+        num_prefix_tokens: Number of virtual prefix tokens to generate. (default: 8)
+    """
+
+    def __init__(
+        self,
+        edge_index: any,
+        edge_type: Optional[any] = None,
+        entity_to_id: Optional[Dict[str, int]] = None,
+        relation_to_id: Optional[Dict[str, int]] = None,
+        x: Optional[any] = None,
+        num_relations: Optional[int] = None,
+        encoder_type: Literal["rgcn", "transe", "none"] = "rgcn",
+        hidden_dim: int = 64,
+        encoder_out_dim: int = 128,
+        llm_dim: int = 4096,
+        num_prefix_tokens: int = 8,
+        connector: Optional[KGLLMConnector] = None,
+    ):
+        self.edge_index = edge_index
+        self.edge_type = edge_type
+        self.x = x
+        self.entity_to_id = entity_to_id or {}
+        self.relation_to_id = relation_to_id or {}
+        self.id_to_entity = {v: k for k, v in self.entity_to_id.items()}
+        self.id_to_relation = {v: k for k, v in self.relation_to_id.items()}
+
+        self.retriever = KGEntityRetriever(
+            entity_to_id=self.entity_to_id,
+            relation_to_id=self.relation_to_id,
+            edge_index=self.edge_index,
+            edge_type=self.edge_type,
+            x=self.x,
+        )
+
+        if connector is not None:
+            self.connector = connector
+        elif encoder_type == "none":
+            self.connector = None
+        else:
+            n_rel = num_relations
+            if n_rel is None and self.relation_to_id:
+                n_rel = len(self.relation_to_id)
+            elif n_rel is None and edge_type is not None:
+                n_rel = int(np.max(ops.convert_to_numpy(edge_type))) + 1
+            else:
+                n_rel = n_rel or 1
+
+            in_ch = ops.shape(x)[-1] if x is not None else hidden_dim
+            if encoder_type == "rgcn":
+                encoder = RGCNSubGraphEncoder(
+                    in_channels=in_ch,
+                    hidden_channels=hidden_dim,
+                    out_channels=encoder_out_dim,
+                    num_relations=n_rel,
+                )
+            elif encoder_type == "transe":
+                num_nodes = len(self.entity_to_id) if self.entity_to_id else int(ops.max(edge_index)) + 1
+                encoder = TransEPrefixEncoder(
+                    num_nodes=num_nodes,
+                    num_relations=n_rel,
+                    embedding_dim=hidden_dim,
+                    out_channels=encoder_out_dim,
+                )
+            else:
+                raise ValueError(f"Unknown encoder_type: {encoder_type}")
+
+            projector = GraphPrefixProjector(
+                in_channels=encoder_out_dim,
+                llm_dim=llm_dim,
+                num_prefix_tokens=num_prefix_tokens,
+            )
+            self.connector = KGLLMConnector(encoder=encoder, projector=projector)
+
+    def retrieve(
+        self,
+        entities: Union[Sequence[Union[str, int]], str, int],
+        num_hops: int = 2,
+        max_nodes_per_hop: Optional[int] = None,
+        directed: bool = False,
+    ) -> SubgraphResult:
+        """Retrieve multi-hop subgraph around specified entities."""
+        return self.retriever.retrieve_subgraph(
+            entities=entities,
+            num_hops=num_hops,
+            max_nodes_per_hop=max_nodes_per_hop,
+            directed=directed,
+        )
+
+    def verbalize(
+        self,
+        subgraph: SubgraphResult,
+        format_style: Literal["triples", "markdown", "natural"] = "markdown",
+        max_triples: Optional[int] = 50,
+    ) -> str:
+        """Verbalize extracted subgraph into text knowledge for LLM prompt context."""
+        return verbalize_subgraph(
+            subgraph=subgraph,
+            id_to_entity=self.id_to_entity,
+            id_to_relation=self.id_to_relation,
+            format_style=format_style,
+            max_triples=max_triples,
+        )
+
+    def build_prompt(
+        self,
+        query: str,
+        entities: Optional[Sequence[Union[str, int]]] = None,
+        num_hops: int = 2,
+        format_style: Literal["triples", "markdown", "natural"] = "markdown",
+        model_family: Literal["llama3", "mistral", "chatml", "standard"] = "llama3",
+        system_prompt: Optional[str] = None,
+    ) -> str:
+        """End-to-end prompt builder: extracts subgraph and formats full LLM prompt."""
+        if entities is None or len(entities) == 0:
+            entities = self.retriever.find_entities_in_text(query)
+
+        subgraph = self.retrieve(entities=entities, num_hops=num_hops)
+        context = self.verbalize(subgraph=subgraph, format_style=format_style)
+        return format_llm_prompt(
+            query=query,
+            context=context,
+            system_prompt=system_prompt,
+            model_family=model_family,
+        )
+
+    def encode_prefix(self, subgraph: SubgraphResult) -> any:
+        """Encode extracted subgraph into LLM prefix embeddings.
+
+        Returns:
+            Tensor of shape `(1, num_prefix_tokens, llm_dim)`.
+        """
+        if self.connector is None:
+            raise ValueError("No KGLLMConnector configured on this GraphRAG instance.")
+        return self.connector.encode_subgraph(subgraph)

k3_node/rag/projector.py

@@ -0,0 +1,184 @@
+"""Prefix Projectors and KG-LLM Connectors for prompt augmentation."""
+
+from typing import Literal, Optional, Tuple, Union
+
+import keras
+from keras import layers, ops
+
+from k3_node.rag.subgraph import SubgraphResult
+
+
+class GraphPrefixProjector(keras.layers.Layer):
+    r"""Projects graph/KGE representations into virtual prefix vectors for LLM prompt augmentation.
+
+    Maps a graph embedding vector of shape `(batch_size, in_channels)` into `num_prefix_tokens`
+    virtual token embeddings of dimension `llm_dim` (e.g. 4096 for Llama 3 / Mistral),
+    suitable for prepending to text token embeddings in LLM forward passes.
+
+    Args:
+        in_channels: Input graph feature dimension from the GNN/KGE encoder.
+        llm_dim: Embedding dimension of the target LLM (e.g., 4096 for Llama-3-8B / Mistral-7B,
+            2048 for Gemma-2B). (default: 4096)
+        num_prefix_tokens: Number of virtual prefix tokens to produce. (default: 8)
+        projector_type: Architecture of the projection head:
+            - `"mlp"`: 2-layer MLP with LayerNorm and GELU activation.
+            - `"linear"`: Single linear transformation.
+        hidden_dim: Optional hidden dimension for MLP projector. (default: `2 * in_channels`)
+        dropout: Dropout probability. (default: 0.0)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        llm_dim: int = 4096,
+        num_prefix_tokens: int = 8,
+        projector_type: Literal["mlp", "linear"] = "mlp",
+        hidden_dim: Optional[int] = None,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.llm_dim = llm_dim
+        self.num_prefix_tokens = num_prefix_tokens
+        self.projector_type = projector_type
+        self.dropout_rate = dropout
+
+        out_total = num_prefix_tokens * llm_dim
+        if hidden_dim is None:
+            hidden_dim = max(in_channels * 2, 512)
+
+        if projector_type == "linear":
+            self.net = layers.Dense(out_total)
+        else:  # "mlp"
+            self.net = keras.Sequential(
+                [
+                    layers.Dense(hidden_dim),
+                    layers.LayerNormalization(),
+                    layers.Activation("gelu"),
+                    layers.Dropout(dropout) if dropout > 0.0 else layers.Identity(),
+                    layers.Dense(out_total),
+                    layers.LayerNormalization(),
+                ]
+            )
+
+    def call(self, graph_embedding: any, training: Optional[bool] = None) -> any:
+        """Project graph embedding into soft prefix token embeddings.
+
+        Args:
+            graph_embedding: Tensor of shape `(batch_size, in_channels)` or `(in_channels,)`.
+
+        Returns:
+            Prefix tensor of shape `(batch_size, num_prefix_tokens, llm_dim)`.
+        """
+        graph_embedding = ops.convert_to_tensor(graph_embedding)
+        # Ensure 2D (batch_size, in_channels)
+        if len(ops.shape(graph_embedding)) == 1:
+            graph_embedding = ops.expand_dims(graph_embedding, axis=0)
+
+        batch_size = ops.shape(graph_embedding)[0]
+        flat_proj = self.net(graph_embedding, training=training)
+        # Reshape to (batch_size, num_prefix_tokens, llm_dim)
+        prefix_tokens = ops.reshape(flat_proj, (batch_size, self.num_prefix_tokens, self.llm_dim))
+        return prefix_tokens
+
+
+class KGLLMConnector(keras.Model):
+    r"""High-level Knowledge Graph to Large Language Model (KG-LLM) Connector.
+
+    Connects a Knowledge Graph encoder (e.g. `RGCNSubGraphEncoder` or `TransEPrefixEncoder`)
+    with a `GraphPrefixProjector` to produce soft prompt prefix embeddings and inject them
+    into Llama 3, Mistral, or other LLMs.
+
+    Example:
+        ```python
+        from k3_node.rag import RGCNSubGraphEncoder, GraphPrefixProjector, KGLLMConnector
+
+        encoder = RGCNSubGraphEncoder(in_channels=16, hidden_channels=32, out_channels=64, num_relations=5)
+        projector = GraphPrefixProjector(in_channels=64, llm_dim=4096, num_prefix_tokens=4)
+        connector = KGLLMConnector(encoder=encoder, projector=projector)
+
+        # Generate prefix embeddings for LLM prompt augmentation
+        prefix_embeds = connector.encode_subgraph(subgraph)  # [1, 4, 4096]
+
+        # Inject into LLM input embeddings [batch, seq_len, 4096]
+        augmented_inputs = connector.inject_prefix(text_token_embeds, prefix_embeds)
+        ```
+
+    Args:
+        encoder: Graph or KGE encoder (e.g., `RGCNSubGraphEncoder`, `TransEPrefixEncoder`).
+        projector: `GraphPrefixProjector` instance.
+    """
+
+    def __init__(
+        self,
+        encoder: keras.layers.Layer,
+        projector: GraphPrefixProjector,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.encoder = encoder
+        self.projector = projector
+
+    @property
+    def num_prefix_tokens(self) -> int:
+        return self.projector.num_prefix_tokens
+
+    @property
+    def llm_dim(self) -> int:
+        return self.projector.llm_dim
+
+    def call(self, *args, **kwargs) -> any:
+        """Encode graph inputs and project directly into LLM prefix tokens."""
+        graph_emb = self.encoder(*args, **kwargs)
+        return self.projector(graph_emb)
+
+    def encode_subgraph(self, subgraph: SubgraphResult) -> any:
+        """Encode an extracted SubgraphResult into LLM prefix tokens.
+
+        Returns:
+            Tensor of shape `(1, num_prefix_tokens, llm_dim)`.
+        """
+        if hasattr(self.encoder, "encode_subgraph"):
+            graph_emb = self.encoder.encode_subgraph(subgraph)
+        else:
+            graph_emb = self.encoder(subgraph.x, subgraph.edge_index)
+        return self.projector(graph_emb)
+
+    @staticmethod
+    def inject_prefix(text_embeddings: any, prefix_embeddings: any) -> any:
+        """Prepend graph prefix embeddings to text token embeddings along sequence dimension.
+
+        Args:
+            text_embeddings: Tensor of shape `(batch, seq_len, llm_dim)`.
+            prefix_embeddings: Tensor of shape `(batch, num_prefix_tokens, llm_dim)`.
+
+        Returns:
+            Concatenated tensor of shape `(batch, num_prefix_tokens + seq_len, llm_dim)`.
+        """
+        text_embeddings = ops.convert_to_tensor(text_embeddings)
+        prefix_embeddings = ops.convert_to_tensor(prefix_embeddings)
+
+        # Match batch size if prefix was computed for single batch
+        b_text = ops.shape(text_embeddings)[0]
+        b_prefix = ops.shape(prefix_embeddings)[0]
+        if b_prefix == 1 and b_text > 1:
+            prefix_embeddings = ops.repeat(prefix_embeddings, b_text, axis=0)
+
+        return ops.concatenate([prefix_embeddings, text_embeddings], axis=1)
+
+    @staticmethod
+    def extend_attention_mask(attention_mask: any, num_prefix_tokens: int) -> any:
+        """Extend LLM binary attention mask with 1s for the prepended prefix tokens.
+
+        Args:
+            attention_mask: Tensor of shape `(batch, seq_len)`.
+            num_prefix_tokens: Number of prefix tokens prepended.
+
+        Returns:
+            Extended mask of shape `(batch, num_prefix_tokens + seq_len)`.
+        """
+        attention_mask = ops.convert_to_tensor(attention_mask)
+        batch_size = ops.shape(attention_mask)[0]
+        prefix_mask = ops.ones((batch_size, num_prefix_tokens), dtype=attention_mask.dtype)
+        return ops.concatenate([prefix_mask, attention_mask], axis=1)