Add token generation for federated server; remove additional P2P setup for federated instance mode

Author: mintyleaf

core/cli/federated.go

@@ -2,10 +2,12 @@ package cli
 
 import (
 	"context"
+	"fmt"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
 	cliP2P "github.com/mudler/LocalAI/core/cli/p2p"
 	"github.com/mudler/LocalAI/core/p2p"
+	"github.com/rs/zerolog/log"
 )
 
 type FederatedCLI struct {
@@ -18,6 +20,24 @@ type FederatedCLI struct {
 
 func (f *FederatedCLI) Run(ctx *cliContext.Context) error {
 
+	if f.Peer2PeerToken == "" {
+		log.Info().Msg("No token provided, generating one")
+		connectionData, err := p2p.GenerateNewConnectionData(
+			f.Peer2PeerDHTInterval, f.Peer2PeerOTPInterval,
+			f.Peer2PeerPrivkey, f.Peer2PeerUsePeerguard,
+		)
+		if err != nil {
+			log.Warn().Msgf("Error generating token: %s", err.Error())
+		}
+		f.Peer2PeerToken = connectionData.Base64()
+
+		log.Info().Msg("Generated Token:")
+		fmt.Println(f.Peer2PeerToken)
+
+		log.Info().Msg("To use the token, you can run the following command in another node or terminal:")
+		fmt.Printf("export TOKEN=\"%s\"\nlocal-ai worker p2p-llama-cpp-rpc\n", f.Peer2PeerToken)
+	}
+
 	fs := p2p.NewFederatedServer(f.Address, p2p.NetworkID(f.Peer2PeerNetworkID, p2p.FederatedID), f.Peer2PeerToken, !f.RandomWorker, f.TargetWorker)
 
 	return fs.Start(context.Background(), f.P2PCommonFlags)

core/cli/run.go

@@ -138,11 +138,6 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 		}
 		opts = append(opts, config.WithP2PToken(token))
 
-		if r.Federated {
-			p2pCfg.PeerGuard.Autocleanup = true
-			p2pCfg.PeerGuard.PeerGate = true
-		}
-
 		p2pCfg.NetworkToken = token
 	}
 

core/p2p/federated_server.go

@@ -15,9 +15,9 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-func (f *FederatedServer) Start(ctx context.Context, p2pCommonFlags cliP2P.P2PCommonFlags) error {
+func (fs *FederatedServer) Start(ctx context.Context, p2pCommonFlags cliP2P.P2PCommonFlags) error {
 	p2pCfg := NewP2PConfig(p2pCommonFlags)
-	p2pCfg.NetworkToken = f.p2ptoken
+	p2pCfg.NetworkToken = fs.p2ptoken
 	p2pCfg.PeerGuard.Autocleanup = true
 	p2pCfg.PeerGuard.PeerGate = true
 
@@ -30,13 +30,13 @@ func (f *FederatedServer) Start(ctx context.Context, p2pCommonFlags cliP2P.P2PCo
 		return fmt.Errorf("creating a new node: %w", err)
 	}
 
-	if err := ServiceDiscoverer(ctx, n, f.service, func(servicesID string, tunnel NodeData) {
+	if err := ServiceDiscoverer(ctx, n, fs.service, func(servicesID string, tunnel NodeData) {
 		log.Debug().Msgf("Discovered node: %s", tunnel.ID)
 	}, false); err != nil {
 		return err
 	}
 
-	return f.proxy(ctx, n)
+	return fs.proxy(ctx, n)
 }
 
 func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {