[RSDK-9677] add GetModelsFromModules to robot interface (viamrobotics…

…#4676)

module/modmanager/manager.go

@@ -612,6 +612,29 @@ func (mgr *Manager) Configs() []config.Module {
 	return configs
 }
 
+// AllModels returns a slice of resource.ModuleModelDiscovery representing the available models
+// from the currently managed modules.
+func (mgr *Manager) AllModels() []resource.ModuleModelDiscovery {
+	moduleTypes := map[string]config.ModuleType{}
+	models := []resource.ModuleModelDiscovery{}
+	for _, moduleConfig := range mgr.Configs() {
+		moduleName := moduleConfig.Name
+		moduleTypes[moduleName] = moduleConfig.Type
+	}
+	for moduleName, handleMap := range mgr.Handles() {
+		for api, handle := range handleMap {
+			for _, model := range handle {
+				modelModel := resource.ModuleModelDiscovery{
+					ModuleName: moduleName, Model: model, API: api.API,
+					FromLocalModule: moduleTypes[moduleName] == config.ModuleTypeLocal,
+				}
+				models = append(models, modelModel)
+			}
+		}
+	}
+	return models
+}
+
 // Provides returns true if a component/service config WOULD be handled by a module.
 func (mgr *Manager) Provides(conf resource.Config) bool {
 	_, ok := mgr.getModule(conf)

module/modmanager/manager_test.go

@@ -206,6 +206,32 @@ func TestModManagerFunctions(t *testing.T) {
 			test.That(t, ok, test.ShouldBeTrue)
 			test.That(t, reg.Constructor, test.ShouldNotBeNil)
 
+			t.Log("test AllModels")
+			modCfg2 := config.Module{
+				Name:    "simple-module2",
+				ExePath: modPath,
+				Type:    config.ModuleTypeLocal,
+			}
+			err = mgr.Add(ctx, modCfg2)
+			test.That(t, err, test.ShouldBeNil)
+			models := mgr.AllModels()
+			for _, model := range models {
+				test.That(t, model.Model, test.ShouldResemble, resource.NewModel("acme", "demo", "mycounter"))
+				test.That(t, model.API, test.ShouldResemble, resource.NewAPI("rdk", "component", "generic"))
+				switch model.ModuleName {
+				case "simple-module":
+					test.That(t, model.FromLocalModule, test.ShouldEqual, false)
+				case "simple-module2":
+					test.That(t, model.FromLocalModule, test.ShouldEqual, true)
+				default:
+					t.Fail()
+					t.Logf("test AllModels failure: unrecoginzed moduleName %v", model.ModuleName)
+				}
+			}
+			names, err := mgr.Remove(modCfg2.Name)
+			test.That(t, names, test.ShouldBeEmpty)
+			test.That(t, err, test.ShouldBeNil)
+
 			t.Log("test Provides")
 			ok = mgr.Provides(cfgCounter1)
 			test.That(t, ok, test.ShouldBeTrue)

module/modmaninterface/interface.go

@@ -24,6 +24,7 @@ type ModuleManager interface {
 	CleanModuleDataDirectory() error
 
 	Configs() []config.Module
+	AllModels() []resource.ModuleModelDiscovery
 	Provides(cfg resource.Config) bool
 	Handles() map[string]module.HandlerMap
 

resource/discovery.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 
+	pb "go.viam.com/api/robot/v1"
+
 	"go.viam.com/rdk/logging"
 )
 
@@ -33,8 +35,24 @@ type (
 		Query DiscoveryQuery
 		Cause error
 	}
+
+	// ModuleModelDiscovery holds the API and Model information of models within a module.
+	ModuleModelDiscovery struct {
+		ModuleName      string
+		API             API
+		Model           Model
+		FromLocalModule bool
+	}
 )
 
+// ToProto converts a ModuleModelDiscovery into the equivalent proto message.
+func (mm *ModuleModelDiscovery) ToProto() *pb.ModuleModel {
+	return &pb.ModuleModel{
+		Model: mm.Model.String(), Api: mm.API.String(), ModuleName: mm.ModuleName,
+		FromLocalModule: mm.FromLocalModule,
+	}
+}
+
 func (e *DiscoverError) Error() string {
 	return fmt.Sprintf("failed to get discovery for api %q and model %q error: %v", e.Query.API, e.Query.Model, e.Cause)
 }

robot/client/client.go

@@ -922,6 +922,32 @@ func (rc *RobotClient) DiscoverComponents(ctx context.Context, qs []resource.Dis
 	return discoveries, nil
 }
 
+// GetModelsFromModules  returns the available models from the configured modules on a given machine.
+func (rc *RobotClient) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) {
+	resp, err := rc.client.GetModelsFromModules(ctx, &pb.GetModelsFromModulesRequest{})
+	if err != nil {
+		return nil, err
+	}
+	protoModels := resp.GetModels()
+	models := []resource.ModuleModelDiscovery{}
+	for _, protoModel := range protoModels {
+		modelTriplet, err := resource.NewModelFromString(protoModel.Model)
+		if err != nil {
+			return nil, err
+		}
+		api, err := resource.NewAPIFromString(protoModel.Api)
+		if err != nil {
+			return nil, err
+		}
+		model := resource.ModuleModelDiscovery{
+			ModuleName: protoModel.ModuleName, Model: modelTriplet, API: api,
+			FromLocalModule: protoModel.FromLocalModule,
+		}
+		models = append(models, model)
+	}
+	return models, nil
+}
+
 // FrameSystemConfig  returns the configuration of the frame system of a given machine.
 //
 //	frameSystem, err := machine.FrameSystemConfig(context.Background(), nil)

robot/client/client_test.go

@@ -1355,6 +1355,60 @@ func TestClientDiscovery(t *testing.T) {
 	test.That(t, err, test.ShouldBeNil)
 }
 
+func TestClientGetModelsFromModules(t *testing.T) {
+	injectRobot := &inject.Robot{}
+	injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil }
+	injectRobot.ResourceNamesFunc = func() []resource.Name {
+		return finalResources
+	}
+	injectRobot.MachineStatusFunc = func(_ context.Context) (robot.MachineStatus, error) {
+		return robot.MachineStatus{State: robot.StateRunning}, nil
+	}
+	expectedModels := []resource.ModuleModelDiscovery{
+		{
+			ModuleName:      "simple-module",
+			API:             resource.NewAPI("rdk", "component", "generic"),
+			Model:           resource.NewModel("acme", "demo", "mycounter"),
+			FromLocalModule: false,
+		},
+		{
+			ModuleName:      "simple-module2",
+			API:             resource.NewAPI("rdk", "component", "generic"),
+			Model:           resource.NewModel("acme", "demo", "mycounter"),
+			FromLocalModule: true,
+		},
+	}
+	injectRobot.GetModelsFromModulesFunc = func(context.Context) ([]resource.ModuleModelDiscovery, error) {
+		return expectedModels, nil
+	}
+
+	gServer := grpc.NewServer()
+	pb.RegisterRobotServiceServer(gServer, server.New(injectRobot))
+	listener, err := net.Listen("tcp", "localhost:0")
+	test.That(t, err, test.ShouldBeNil)
+	logger := logging.NewTestLogger(t)
+
+	go gServer.Serve(listener)
+	defer gServer.Stop()
+
+	client, err := New(context.Background(), listener.Addr().String(), logger)
+	test.That(t, err, test.ShouldBeNil)
+
+	resp, err := client.GetModelsFromModules(context.Background())
+	test.That(t, err, test.ShouldBeNil)
+	test.That(t, len(resp), test.ShouldEqual, 2)
+	test.That(t, resp, test.ShouldResemble, expectedModels)
+	for index, model := range resp {
+		test.That(t, model.ModuleName, test.ShouldEqual, expectedModels[index].ModuleName)
+		test.That(t, model.Model, test.ShouldResemble, expectedModels[index].Model)
+		test.That(t, model.API, test.ShouldResemble, expectedModels[index].API)
+		test.That(t, model.FromLocalModule, test.ShouldEqual, expectedModels[index].FromLocalModule)
+	}
+
+	err = client.Close(context.Background())
+	test.That(t, err, test.ShouldBeNil)
+}
+
 func ensurePartsAreEqual(part, otherPart *referenceframe.FrameSystemPart) error {
 	if part.FrameConfig.Name() != otherPart.FrameConfig.Name() {
 		return fmt.Errorf("part had name %s while other part had name %s", part.FrameConfig.Name(), otherPart.FrameConfig.Name())

robot/impl/discovery_test.go

@@ -8,10 +8,19 @@ import (
 	modulepb "go.viam.com/api/module/v1"
 	"go.viam.com/test"
 
+	"go.viam.com/rdk/components/base"
+	"go.viam.com/rdk/components/generic"
 	"go.viam.com/rdk/config"
+	"go.viam.com/rdk/examples/customresources/apis/gizmoapi"
+	"go.viam.com/rdk/examples/customresources/apis/summationapi"
+	"go.viam.com/rdk/examples/customresources/models/mybase"
+	"go.viam.com/rdk/examples/customresources/models/mygizmo"
+	"go.viam.com/rdk/examples/customresources/models/mynavigation"
+	"go.viam.com/rdk/examples/customresources/models/mysum"
 	"go.viam.com/rdk/logging"
 	"go.viam.com/rdk/resource"
 	"go.viam.com/rdk/robot"
+	"go.viam.com/rdk/services/navigation"
 	rtestutils "go.viam.com/rdk/testutils"
 )
 
@@ -168,3 +177,66 @@ func TestDiscovery(t *testing.T) {
 		test.That(t, len(complexHandles.Handlers), test.ShouldBeGreaterThan, 1)
 	})
 }
+
+func TestGetModelsFromModules(t *testing.T) {
+	t.Run("no modules configured", func(t *testing.T) {
+		r := setupLocalRobotWithFakeConfig(t)
+		models, err := r.GetModelsFromModules(context.Background())
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, models, test.ShouldBeEmpty)
+	})
+	t.Run("local and registry modules are configured", func(t *testing.T) {
+		r := setupLocalRobotWithFakeConfig(t)
+		ctx := context.Background()
+
+		// add modules
+		complexPath := rtestutils.BuildTempModule(t, "examples/customresources/demos/complexmodule")
+		simplePath := rtestutils.BuildTempModule(t, "examples/customresources/demos/simplemodule")
+		cfg := &config.Config{
+			Modules: []config.Module{
+				{
+					Name:    "simple",
+					ExePath: simplePath,
+					Type:    config.ModuleTypeRegistry,
+				},
+				{
+					Name:    "complex",
+					ExePath: complexPath,
+					Type:    config.ModuleTypeLocal,
+				},
+			},
+		}
+		r.Reconfigure(ctx, cfg)
+		models, err := r.GetModelsFromModules(context.Background())
+		test.That(t, err, test.ShouldBeNil)
+		test.That(t, models, test.ShouldHaveLength, 5)
+
+		for _, model := range models {
+			switch model.Model {
+			case mygizmo.Model:
+				test.That(t, model.FromLocalModule, test.ShouldEqual, true)
+				test.That(t, model.ModuleName, test.ShouldEqual, "complex")
+				test.That(t, model.API, test.ShouldResemble, gizmoapi.API)
+			case mysum.Model:
+				test.That(t, model.FromLocalModule, test.ShouldEqual, true)
+				test.That(t, model.ModuleName, test.ShouldEqual, "complex")
+				test.That(t, model.API, test.ShouldResemble, summationapi.API)
+			case mybase.Model:
+				test.That(t, model.FromLocalModule, test.ShouldEqual, true)
+				test.That(t, model.ModuleName, test.ShouldEqual, "complex")
+				test.That(t, model.API, test.ShouldResemble, base.API)
+			case mynavigation.Model:
+				test.That(t, model.FromLocalModule, test.ShouldEqual, true)
+				test.That(t, model.ModuleName, test.ShouldEqual, "complex")
+				test.That(t, model.API, test.ShouldResemble, navigation.API)
+			case resource.NewModel("acme", "demo", "mycounter"):
+				test.That(t, model.FromLocalModule, test.ShouldEqual, false)
+				test.That(t, model.ModuleName, test.ShouldEqual, "simple")
+				test.That(t, model.API, test.ShouldResemble, generic.API)
+			default:
+				t.Fail()
+				t.Logf("test GetModelsFromModules failure: unrecoginzed model %v", model.Model)
+			}
+		}
+	})
+}

robot/impl/local_robot.go

@@ -1077,6 +1077,10 @@ func (r *localRobot) discoverRobotInternals(query resource.DiscoveryQuery) (inte
 	}
 }
 
+func (r *localRobot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) {
+	return r.manager.moduleManager.AllModels(), nil
+}
+
 func dialRobotClient(
 	ctx context.Context,
 	config config.Remote,

robot/impl/resource_manager_test.go

@@ -1886,6 +1886,15 @@ func (rr *dummyRobot) DiscoverComponents(ctx context.Context, qs []resource.Disc
 	return rr.robot.DiscoverComponents(ctx, qs)
 }
 
+func (rr *dummyRobot) GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error) {
+	rr.mu.Lock()
+	defer rr.mu.Unlock()
+	if rr.offline {
+		return nil, errors.New("offline")
+	}
+	return rr.robot.GetModelsFromModules(ctx)
+}
+
 func (rr *dummyRobot) RemoteNames() []string {
 	return nil
 }

robot/robot.go

@@ -88,6 +88,10 @@ type Robot interface {
 	// Only implemented for webcam cameras in builtin components.
 	DiscoverComponents(ctx context.Context, qs []resource.DiscoveryQuery) ([]resource.Discovery, error)
 
+	// GetModelsFromModules returns a list of models supported by the configured modules,
+	// and specifies whether the models are from a local or registry module.
+	GetModelsFromModules(ctx context.Context) ([]resource.ModuleModelDiscovery, error)
+
 	// RemoteByName returns a remote robot by name.
 	RemoteByName(name string) (Robot, bool)
 

robot/server/server.go

@@ -228,6 +228,19 @@ func (s *Server) DiscoverComponents(ctx context.Context, req *pb.DiscoverCompone
 	return &pb.DiscoverComponentsResponse{Discovery: pbDiscoveries}, nil
 }
 
+// GetModelsFromModules returns all models from the currently managed modules.
+func (s *Server) GetModelsFromModules(ctx context.Context, req *pb.GetModelsFromModulesRequest) (*pb.GetModelsFromModulesResponse, error) {
+	models, err := s.robot.GetModelsFromModules(ctx)
+	if err != nil {
+		return nil, err
+	}
+	resp := pb.GetModelsFromModulesResponse{}
+	for _, mm := range models {
+		resp.Models = append(resp.Models, mm.ToProto())
+	}
+	return &resp, nil
+}
+
 // FrameSystemConfig returns the info of each individual part that makes up the frame system.
 func (s *Server) FrameSystemConfig(ctx context.Context, req *pb.FrameSystemConfigRequest) (*pb.FrameSystemConfigResponse, error) {
 	fsCfg, err := s.robot.FrameSystemConfig(ctx)

robot/server/server_test.go

@@ -448,6 +448,49 @@ func TestServer(t *testing.T) {
 		})
 	})
 
+	t.Run("GetModelsFromModules", func(t *testing.T) {
+		injectRobot := &inject.Robot{}
+		injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil }
+		injectRobot.ResourceNamesFunc = func() []resource.Name { return []resource.Name{} }
+		server := server.New(injectRobot)
+
+		expectedModels := []resource.ModuleModelDiscovery{
+			{
+				ModuleName:      "simple-module",
+				API:             resource.NewAPI("rdk", "component", "generic"),
+				Model:           resource.NewModel("acme", "demo", "mycounter"),
+				FromLocalModule: false,
+			},
+			{
+				ModuleName:      "simple-module2",
+				API:             resource.NewAPI("rdk", "component", "generic"),
+				Model:           resource.NewModel("acme", "demo", "mycounter"),
+				FromLocalModule: true,
+			},
+		}
+		injectRobot.GetModelsFromModulesFunc = func(context.Context) ([]resource.ModuleModelDiscovery, error) {
+			return expectedModels, nil
+		}
+		expectedProto := []*pb.ModuleModel{expectedModels[0].ToProto(), expectedModels[1].ToProto()}
+
+		req := &pb.GetModelsFromModulesRequest{}
+		resp, err := server.GetModelsFromModules(context.Background(), req)
+		test.That(t, err, test.ShouldBeNil)
+		protoModels := resp.GetModels()
+		test.That(t, len(protoModels), test.ShouldEqual, 2)
+		test.That(t, protoModels, test.ShouldResemble, expectedProto)
+		for index, protoModel := range protoModels {
+			test.That(t, protoModel.ModuleName, test.ShouldEqual, expectedProto[index].ModuleName)
+			test.That(t, protoModel.ModuleName, test.ShouldEqual, expectedModels[index].ModuleName)
+			test.That(t, protoModel.Model, test.ShouldEqual, expectedProto[index].Model)
+			test.That(t, protoModel.Model, test.ShouldEqual, expectedModels[index].Model.String())
+			test.That(t, protoModel.Api, test.ShouldEqual, expectedProto[index].Api)
+			test.That(t, protoModel.Api, test.ShouldEqual, expectedModels[index].API.String())
+			test.That(t, protoModel.FromLocalModule, test.ShouldEqual, expectedProto[index].FromLocalModule)
+			test.That(t, protoModel.FromLocalModule, test.ShouldEqual, expectedModels[index].FromLocalModule)
+		}
+	})
+
 	t.Run("ResourceRPCSubtypes", func(t *testing.T) {
 		injectRobot := &inject.Robot{}
 		injectRobot.ResourceRPCAPIsFunc = func() []resource.RPCAPI { return nil }