diff --git a/module/modmanager/manager.go b/module/modmanager/manager.go index 194532c4fad..0ae6de23387 100644 --- a/module/modmanager/manager.go +++ b/module/modmanager/manager.go @@ -596,6 +596,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) diff --git a/module/modmaninterface/interface.go b/module/modmaninterface/interface.go index 22f589ccee4..411880657da 100644 --- a/module/modmaninterface/interface.go +++ b/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 diff --git a/resource/discovery.go b/resource/discovery.go index 9b2a11f6cba..f4b6c8e686b 100644 --- a/resource/discovery.go +++ b/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) } diff --git a/robot/client/client.go b/robot/client/client.go index 32a641052a2..189f87b2a2a 100644 --- a/robot/client/client.go +++ b/robot/client/client.go @@ -876,6 +876,34 @@ 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. +// +// models, err := machine.GetModelsFromModules(context.Background()) +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) diff --git a/robot/impl/local_robot.go b/robot/impl/local_robot.go index 5189618a5fe..577636e78bb 100644 --- a/robot/impl/local_robot.go +++ b/robot/impl/local_robot.go @@ -1079,6 +1079,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, diff --git a/robot/impl/resource_manager_test.go b/robot/impl/resource_manager_test.go index 4a2e9c39c7a..2b721bea120 100644 --- a/robot/impl/resource_manager_test.go +++ b/robot/impl/resource_manager_test.go @@ -1894,6 +1894,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 } diff --git a/robot/robot.go b/robot/robot.go index f57330d0852..64482b63af6 100644 --- a/robot/robot.go +++ b/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) diff --git a/robot/server/server.go b/robot/server/server.go index 86c3df4a8ee..c23e250bc40 100644 --- a/robot/server/server.go +++ b/robot/server/server.go @@ -223,6 +223,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)