rag/internal/cmd/rpc.go
2024-07-15 21:08:48 +08:00

77 lines
2.0 KiB
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)
}