torch_geometric.llm.utils.NeighborSamplingRAGGraphStore

class NeighborSamplingRAGGraphStore(feature_store: Optional[FeatureStore] = None, **kwargs)[source]

Bases: LocalGraphStore

Neighbor sampling based graph-store to store & retrieve graph data.

property config: Dict[str, Any]

Get the config for the feature store.

Return type:

Dict[str, Any]

register_feature_store(feature_store: FeatureStore) → None[source]

Registers a feature store with the graph store.

Parameters:

feature_store (FeatureStore) – The feature store to register.

Return type:

None

put_edge_id(edge_id: Tensor, *args, **kwargs) → bool[source]

Stores an edge ID in the graph store.

Parameters:

edge_id (Tensor) – The edge ID to store.

Returns:

bool – Whether the operation was successful.

put_edge_index(edge_index: Union[Tensor, Tuple[Tensor, Tensor]], *args, **kwargs) → bool[source]

Stores an edge index in the graph store.

Parameters:

edge_index (Union[Tensor, Tuple[Tensor, Tensor]]) – The edge index to store.

Returns:

bool – Whether the operation was successful.

property edge_index: Union[Tensor, Tuple[Tensor, Tensor]]

Gets the edge index of the graph.

Returns:

Union[Tensor, Tuple[Tensor, Tensor]] – The edge index as a tensor.

sample_subgraph(seed_nodes: Tensor) → SamplerOutput[source]

Sample the graph starting from the given nodes using the in-built NeighborSampler.

Parameters:

  • seed_nodes (InputNodes) – Seed nodes to start sampling from.

  • num_neighbors (Optional[NumNeighborsType], optional) – Parameters to determine how many hops and number of neighbors per hop. Defaults to None.

Returns:

NeighborSamplerOutput

for the input.

Return type:

Union[SamplerOutput, HeteroSamplerOutput]