feat(bedrockagent): add agent, prompt, datasource and knowledge base

resources/bedrock-agent-datasource.go

@@ -0,0 +1,114 @@
+package resources
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/bedrockagent"
+
+	"github.com/ekristen/libnuke/pkg/registry"
+	"github.com/ekristen/libnuke/pkg/resource"
+	"github.com/ekristen/libnuke/pkg/types"
+
+	"github.com/ekristen/aws-nuke/v3/pkg/nuke"
+)
+
+const BedrockDataSourceResource = "BedrockDataSource"
+
+func init() {
+	registry.Register(&registry.Registration{
+		Name:   BedrockDataSourceResource,
+		Scope:  nuke.Account,
+		Lister: &BedrockDataSourceLister{},
+	})
+}
+
+type BedrockDataSourceLister struct{}
+
+func (l *BedrockDataSourceLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
+	opts := o.(*nuke.ListerOpts)
+
+	svc := bedrockagent.New(opts.Session)
+	resources := make([]resource.Resource, 0)
+
+	// Require to have workspaces identifiers to query components
+	knowledgeBaseListResponse, err := ListKnowledgeBaseForDataSource(svc)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, knowledgeBaseResponse := range knowledgeBaseListResponse {
+		params := &bedrockagent.ListDataSourcesInput{
+			KnowledgeBaseId: knowledgeBaseResponse.KnowledgeBaseId,
+			MaxResults:      aws.Int64(25),
+		}
+
+		for {
+			resp, err := svc.ListDataSources(params)
+			if err != nil {
+				return nil, err
+			}
+
+			for _, item := range resp.DataSourceSummaries {
+				resources = append(resources, &BedrockDataSource{
+					svc:             svc,
+					ID:              item.DataSourceId,
+					Name:            item.Name,
+					Status:          item.Status,
+					KnowledgeBaseID: knowledgeBaseResponse.KnowledgeBaseId,
+				})
+			}
+			if resp.NextToken == nil {
+				break
+			}
+			params.NextToken = resp.NextToken
+		}
+	}
+
+	return resources, nil
+}
+
+// Utility function to list knowledge bases
+func ListKnowledgeBaseForDataSource(svc *bedrockagent.BedrockAgent) ([]*bedrockagent.KnowledgeBaseSummary, error) {
+	resources := make([]*bedrockagent.KnowledgeBaseSummary, 0)
+	params := &bedrockagent.ListKnowledgeBasesInput{
+		MaxResults: aws.Int64(25),
+	}
+	for {
+		resp, err := svc.ListKnowledgeBases(params)
+		if err != nil {
+			return nil, err
+		}
+		resources = append(resources, resp.KnowledgeBaseSummaries...)
+		if resp.NextToken == nil {
+			break
+		}
+		params.NextToken = resp.NextToken
+	}
+	return resources, nil
+}
+
+type BedrockDataSource struct {
+	svc             *bedrockagent.BedrockAgent
+	ID              *string
+	Name            *string
+	Status          *string
+	KnowledgeBaseID *string
+}
+
+func (r *BedrockDataSource) Properties() types.Properties {
+	return types.NewPropertiesFromStruct(r)
+}
+
+func (r *BedrockDataSource) Remove(_ context.Context) error {
+	_, err := r.svc.DeleteDataSource(&bedrockagent.DeleteDataSourceInput{
+		DataSourceId:    r.ID,
+		KnowledgeBaseId: r.KnowledgeBaseID,
+	})
+
+	return err
+}
+
+func (r *BedrockDataSource) String() string {
+	return *r.ID
+}

resources/bedrock-agent-knowledge-base.go

@@ -0,0 +1,86 @@
+package resources
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/bedrockagent"
+
+	"github.com/ekristen/libnuke/pkg/registry"
+	"github.com/ekristen/libnuke/pkg/resource"
+	"github.com/ekristen/libnuke/pkg/types"
+
+	"github.com/ekristen/aws-nuke/v3/pkg/nuke"
+)
+
+const BedrockKnowledgeBaseResource = "BedrockKnowledgeBase"
+
+func init() {
+	registry.Register(&registry.Registration{
+		Name:   BedrockKnowledgeBaseResource,
+		Scope:  nuke.Account,
+		Lister: &BedrockKnowledgeBaseLister{},
+		DependsOn: []string{
+			BedrockDataSourceResource,
+		},
+	})
+}
+
+type BedrockKnowledgeBaseLister struct{}
+
+func (l *BedrockKnowledgeBaseLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
+	opts := o.(*nuke.ListerOpts)
+
+	svc := bedrockagent.New(opts.Session)
+	resources := make([]resource.Resource, 0)
+
+	params := &bedrockagent.ListKnowledgeBasesInput{
+		MaxResults: aws.Int64(25),
+	}
+
+	for {
+		resp, err := svc.ListKnowledgeBases(params)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, item := range resp.KnowledgeBaseSummaries {
+			resources = append(resources, &BedrockKnowledgeBase{
+				svc:    svc,
+				ID:     item.KnowledgeBaseId,
+				Name:   item.Name,
+				Status: item.Status,
+			})
+		}
+		if resp.NextToken == nil {
+			break
+		}
+
+		params.NextToken = resp.NextToken
+	}
+
+	return resources, nil
+}
+
+type BedrockKnowledgeBase struct {
+	svc    *bedrockagent.BedrockAgent
+	ID     *string
+	Name   *string
+	Status *string
+}
+
+func (r *BedrockKnowledgeBase) Properties() types.Properties {
+	return types.NewPropertiesFromStruct(r)
+}
+
+func (r *BedrockKnowledgeBase) Remove(_ context.Context) error {
+	_, err := r.svc.DeleteKnowledgeBase(&bedrockagent.DeleteKnowledgeBaseInput{
+		KnowledgeBaseId: r.ID,
+	})
+
+	return err
+}
+
+func (r *BedrockKnowledgeBase) String() string {
+	return *r.ID
+}

resources/bedrock-agent-prompt.go

@@ -0,0 +1,84 @@
+package resources
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/bedrockagent"
+
+	"github.com/ekristen/libnuke/pkg/registry"
+	"github.com/ekristen/libnuke/pkg/resource"
+	"github.com/ekristen/libnuke/pkg/types"
+
+	"github.com/ekristen/aws-nuke/v3/pkg/nuke"
+)
+
+const BedrockPromptResource = "BedrockPrompt"
+
+func init() {
+	registry.Register(&registry.Registration{
+		Name:   BedrockPromptResource,
+		Scope:  nuke.Account,
+		Lister: &BedrockPromptLister{},
+	})
+}
+
+type BedrockPromptLister struct{}
+
+func (l *BedrockPromptLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
+	opts := o.(*nuke.ListerOpts)
+
+	svc := bedrockagent.New(opts.Session)
+	resources := make([]resource.Resource, 0)
+
+	params := &bedrockagent.ListPromptsInput{
+		MaxResults: aws.Int64(25),
+	}
+
+	for {
+		resp, err := svc.ListPrompts(params)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, item := range resp.PromptSummaries {
+			resources = append(resources, &BedrockPrompt{
+				svc:     svc,
+				ID:      item.Id,
+				Name:    item.Name,
+				Version: item.Version,
+			})
+		}
+		if resp.NextToken == nil {
+			break
+		}
+
+		params.NextToken = resp.NextToken
+	}
+
+	return resources, nil
+}
+
+type BedrockPrompt struct {
+	svc     *bedrockagent.BedrockAgent
+	ID      *string
+	Name    *string
+	Version *string
+}
+
+func (r *BedrockPrompt) Properties() types.Properties {
+	return types.NewPropertiesFromStruct(r)
+}
+
+func (r *BedrockPrompt) Remove(_ context.Context) error {
+	// Without PromptVersion param, will delete the prompt and all its versions
+	_, err := r.svc.DeletePrompt(&bedrockagent.DeletePromptInput{
+		PromptIdentifier: r.ID,
+	})
+
+	return err
+}
+
+func (r *BedrockPrompt) String() string {
+	return *r.ID
+}

resources/bedrock-agent.go

@@ -0,0 +1,86 @@
+package resources
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/bedrockagent"
+
+	"github.com/ekristen/libnuke/pkg/registry"
+	"github.com/ekristen/libnuke/pkg/resource"
+	"github.com/ekristen/libnuke/pkg/types"
+
+	"github.com/ekristen/aws-nuke/v3/pkg/nuke"
+)
+
+const BedrockAgentResource = "BedrockAgent"
+
+func init() {
+	registry.Register(&registry.Registration{
+		Name:   BedrockAgentResource,
+		Scope:  nuke.Account,
+		Lister: &BedrockAgentLister{},
+	})
+}
+
+type BedrockAgentLister struct{}
+
+func (l *BedrockAgentLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
+	opts := o.(*nuke.ListerOpts)
+
+	svc := bedrockagent.New(opts.Session)
+	resources := make([]resource.Resource, 0)
+
+	params := &bedrockagent.ListAgentsInput{
+		MaxResults: aws.Int64(25),
+	}
+
+	for {
+		resp, err := svc.ListAgents(params)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, item := range resp.AgentSummaries {
+			// Cannot query tags directly here, AgentSummaries do not contain agent ARN...
+
+			resources = append(resources, &BedrockAgent{
+				svc:    svc,
+				ID:     item.AgentId,
+				Name:   item.AgentName,
+				Status: item.AgentStatus,
+			})
+		}
+		if resp.NextToken == nil {
+			break
+		}
+
+		params.NextToken = resp.NextToken
+	}
+
+	return resources, nil
+}
+
+type BedrockAgent struct {
+	svc    *bedrockagent.BedrockAgent
+	ID     *string
+	Name   *string
+	Status *string
+}
+
+func (r *BedrockAgent) Properties() types.Properties {
+	return types.NewPropertiesFromStruct(r)
+}
+
+func (r *BedrockAgent) Remove(_ context.Context) error {
+	_, err := r.svc.DeleteAgent(&bedrockagent.DeleteAgentInput{
+		AgentId:                r.ID,
+		SkipResourceInUseCheck: aws.Bool(true),
+	})
+
+	return err
+}
+
+func (r *BedrockAgent) String() string {
+	return *r.ID
+}