Skip to content

Commit

Permalink
feat(bedrock): add guardrails and evaluation jobs
Browse files Browse the repository at this point in the history
  • Loading branch information
npellegrin committed Nov 18, 2024
1 parent 674525d commit 9bd88a5
Show file tree
Hide file tree
Showing 2 changed files with 202 additions and 0 deletions.
105 changes: 105 additions & 0 deletions resources/bedrock-evaluation-jobs.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
package resources

import (
"context"
"fmt"

"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrock"

"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 BedrockEvaluationJobResource = "BedrockEvaluationJob"

func init() {
registry.Register(&registry.Registration{
Name: BedrockEvaluationJobResource,
Scope: nuke.Account,
Lister: &BedrockEvaluationJobLister{},
})
}

type BedrockEvaluationJobLister struct{}

func (l *BedrockEvaluationJobLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
opts := o.(*nuke.ListerOpts)

svc := bedrock.New(opts.Session)
resources := make([]resource.Resource, 0)

params := &bedrock.ListEvaluationJobsInput{
MaxResults: aws.Int64(30),
}

for {
resp, err := svc.ListEvaluationJobs(params)
if err != nil {
return nil, err
}

for _, jobSummary := range resp.JobSummaries {
tagResp, err := svc.ListTagsForResource(
&bedrock.ListTagsForResourceInput{
ResourceARN: jobSummary.JobArn,
})
if err != nil {
return nil, err
}

resources = append(resources, &BedrockEvaluationJob{
svc: svc,
Arn: jobSummary.JobArn,
Name: jobSummary.JobName,
Status: jobSummary.Status,
Tags: tagResp.Tags,
})
}

if resp.NextToken == nil {
break
}

params.NextToken = resp.NextToken
}

return resources, nil
}

type BedrockEvaluationJob struct {
svc *bedrock.Bedrock
Arn *string
Name *string
Status *string
Tags []*bedrock.Tag
}

func (r *BedrockEvaluationJob) Remove(_ context.Context) error {
// We cannot delete an evaluation job from API, only stop it
// Deletion seems to be possible in console only for now
_, err := r.svc.StopEvaluationJob(&bedrock.StopEvaluationJobInput{
JobIdentifier: r.Arn,
})

return err
}

func (r *BedrockEvaluationJob) String() string {
return *r.Name
}

func (r *BedrockEvaluationJob) Properties() types.Properties {
return types.NewPropertiesFromStruct(r)
}

func (r *BedrockEvaluationJob) Filter() error {
if *r.Status != bedrock.EvaluationJobStatusInProgress {
// May be completed, failed, stopping or stopped
return fmt.Errorf("already stopped")
}
return nil
}
97 changes: 97 additions & 0 deletions resources/bedrock-guardrails.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
package resources

import (
"context"

"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrock"

"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 BedrockGuardrailResource = "BedrockGuardrail"

func init() {
registry.Register(&registry.Registration{
Name: BedrockGuardrailResource,
Scope: nuke.Account,
Lister: &BedrockGuardrailLister{},
})
}

type BedrockGuardrailLister struct{}

func (l *BedrockGuardrailLister) List(_ context.Context, o interface{}) ([]resource.Resource, error) {
opts := o.(*nuke.ListerOpts)

svc := bedrock.New(opts.Session)
resources := make([]resource.Resource, 0)

params := &bedrock.ListGuardrailsInput{
MaxResults: aws.Int64(30),
}

for {
resp, err := svc.ListGuardrails(params)
if err != nil {
return nil, err
}

for _, guardrail := range resp.Guardrails {
tagResp, err := svc.ListTagsForResource(
&bedrock.ListTagsForResourceInput{
ResourceARN: guardrail.Arn,
})
if err != nil {
return nil, err
}

resources = append(resources, &BedrockGuardrail{
svc: svc,
ID: guardrail.Id,
Version: guardrail.Version,
Name: guardrail.Name,
Status: guardrail.Status,
Tags: tagResp.Tags,
})
}

if resp.NextToken == nil {
break
}

params.NextToken = resp.NextToken
}

return resources, nil
}

type BedrockGuardrail struct {
svc *bedrock.Bedrock
ID *string
Version *string
Name *string
Status *string
Tags []*bedrock.Tag
}

func (r *BedrockGuardrail) String() string {
return *r.ID
}

func (r *BedrockGuardrail) Remove(_ context.Context) error {
// When guardrail version is not specified, all versions are deleted
_, err := r.svc.DeleteGuardrail(&bedrock.DeleteGuardrailInput{
GuardrailIdentifier: r.ID,
})

return err
}

func (r *BedrockGuardrail) Properties() types.Properties {
return types.NewPropertiesFromStruct(r)
}

0 comments on commit 9bd88a5

Please sign in to comment.