From 3a707e51ef252c393a20fe2f6e446c8d9ed78dfb Mon Sep 17 00:00:00 2001 From: Nicolas PELLEGRIN Date: Fri, 22 Nov 2024 00:43:21 +0100 Subject: [PATCH] feat(bedrockagent): add agent, prompt, datasource and knowledge base --- resources/bedrock-agent-datasource.go | 114 ++++++++++++++++++++++ resources/bedrock-agent-knowledge-base.go | 86 ++++++++++++++++ resources/bedrock-agent-prompt.go | 84 ++++++++++++++++ resources/bedrock-agent.go | 86 ++++++++++++++++ 4 files changed, 370 insertions(+) create mode 100644 resources/bedrock-agent-datasource.go create mode 100644 resources/bedrock-agent-knowledge-base.go create mode 100644 resources/bedrock-agent-prompt.go create mode 100644 resources/bedrock-agent.go diff --git a/resources/bedrock-agent-datasource.go b/resources/bedrock-agent-datasource.go new file mode 100644 index 00000000..8f73eadb --- /dev/null +++ b/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(®istry.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 +} diff --git a/resources/bedrock-agent-knowledge-base.go b/resources/bedrock-agent-knowledge-base.go new file mode 100644 index 00000000..8e3f7ad3 --- /dev/null +++ b/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(®istry.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 +} diff --git a/resources/bedrock-agent-prompt.go b/resources/bedrock-agent-prompt.go new file mode 100644 index 00000000..da233592 --- /dev/null +++ b/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(®istry.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 +} diff --git a/resources/bedrock-agent.go b/resources/bedrock-agent.go new file mode 100644 index 00000000..49311244 --- /dev/null +++ b/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(®istry.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 +}