rag/internal/cmd/rpc.go

package cmd

import (
	"context"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/auth"
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
	"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/reflection"
	"leafdev.top/leaf/rag/api/library"
	libraryGw "leafdev.top/leaf/rag/api/library"
	grpc2 "leafdev.top/leaf/rag/internal/middleware/grpc"
	"leafdev.top/leaf/rag/internal/providers/jwks"
	"leafdev.top/leaf/rag/internal/services/libraryService"
	"net"
)

var rpcCommand = &cobra.Command{
	Use: "rpc",

	Run: func(cmd *cobra.Command, args []string) {
		jwks.InitJwksRefresh()
		StartGRPC()
	},
}

func StartGRPC() {
	if config.ListenAddr.GRPC == "" {
		config.ListenAddr.GRPC = "0.0.0.0:8081"
	}

	lis, err := net.Listen("tcp", config.ListenAddr.GRPC)
	if err != nil {
		panic("failed to listen: " + err.Error())
	}
	logger.Info("Server listening at " + config.ListenAddr.GRPC)

	var opts = []grpc.ServerOption{
		grpc.ChainUnaryInterceptor(
			logging.UnaryServerInterceptor(grpc2.ZapLogInterceptor()),
			auth.UnaryServerInterceptor(grpc2.JwtAuth),
		),
		grpc.ChainStreamInterceptor(
			logging.StreamServerInterceptor(grpc2.ZapLogInterceptor()),
			auth.StreamServerInterceptor(grpc2.JwtAuth),
		),
	}
	grpcServer := grpc.NewServer(opts...)

	reflection.Register(grpcServer)

	library.RegisterLibraryServiceServer(grpcServer, libraryService.LibraryService{})

	err = grpcServer.Serve(lis)
	if err != nil {
		panic(err)
		return
	}

	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	mux := runtime.NewServeMux()
	clientOpts := []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
	err := gw.RegisterYourServiceHandlerFromEndpoint(ctx, mux, config.ListenAddr.GRPC, clientOpts)
	if err != nil {
		return err
	}

	// Start HTTP server (and proxy calls to gRPC server endpoint)
	return http.ListenAndServe(":8081", mux)

}