code: def create_graphSAGE_model(graph, link_prediction=False):
      if link_prediction:
          # We are going to train on the original graph
          generator = GraphSAGELinkGenerator(
              graph, batch_size=2, num_samples=[2, 2]
          )
          edge_ids_train = np.array([[1, 2], [2, 3], [1, 3]])
          train_gen = generator.flow(
              edge_ids_train, np.array([1, 1, 0])
          )
      else:
          generator = GraphSAGENodeGenerator(
              graph, batch_size=2, num_samples=[2, 2]
          )
          train_gen = generator.flow([1, 2], np.array([[1, 0], [0, 1]]))
  
      # if link_prediction:
      #     edge_ids_train = np.array([[1, 2], [2, 3], [1, 3]])
      #     train_gen = generator.flow(edge_ids_train, np.array([1, 1, 0]))
      # else:
      #     train_gen = generator.flow([1, 2], np.array([[1, 0], [0, 1]]))
  
      base_model = GraphSAGE(
          layer_sizes=[8, 8],
          generator=generator,
          bias=True,
          dropout=0.5,
      )
  
      if link_prediction:
          # Expose input and output sockets of graphsage, for source and destination nodes:
          x_inp, x_out = base_model.in_out_tensors()
  
          prediction = link_classification(
              output_dim=1,
              output_act="relu",
              edge_embedding_method="ip",
          )(x_out)
  
          keras_model = Model(inputs=x_inp, outputs=prediction)
      else:
          x_inp, x_out = base_model.in_out_tensors()
          prediction = layers.Dense(units=2, activation="softmax")(
              x_out
          )
  
          keras_model = Model(inputs=x_inp, outputs=prediction)
  
      return base_model, keras_model, generator, train_gen
  
imports: import pandas
  import pytest
  import numpy
  import tensorflow
  from stellargraph import StellarGraph
  from stellargraph.layer import GraphSAGE, GCN, GAT, HinSAGE, link_classification, link_regression
  from stellargraph.mapper import GraphSAGENodeGenerator, FullBatchNodeGenerator, HinSAGENodeGenerator, GraphSAGELinkGenerator, HinSAGELinkGenerator
  from stellargraph.ensemble import Ensemble, BaggingEnsemble
  from tensorflow.keras import layers, Model
  from tensorflow.keras.optimizers import Adam
  from tensorflow.keras.losses import categorical_crossentropy, binary_crossentropy
label_desc: Use image node_id 2 for calling a global function with example usage: create_graphSAGE_model(graph, link_prediction) and returns: base_model,  keras_model,  generator,  train_gen