Skip to content
This repository was archived by the owner on Apr 14, 2025. It is now read-only.

Commit 9021b43

Browse files
authored
Merge pull request #20 from samalba/dockerfile-optimzer2
docker-optimizer: complete rewrite with a smaller prompt and smaller toolset
2 parents 2b1b0b3 + 05135ca commit 9021b43

8 files changed

Lines changed: 180 additions & 90 deletions

File tree

dockerfile-optimizer/dagger.json

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
11
{
22
"name": "dockerfile-optimizer",
3-
"engineVersion": "v0.17.0-llm.3",
3+
"engineVersion": "v0.17.0-llm.8",
44
"sdk": {
55
"source": "go"
66
},
77
"dependencies": [
8+
{
9+
"name": "feature-branch",
10+
"source": "github.com/kpenfound/dag/feature-branch",
11+
"pin": "d1c42dcd8f29d44370286e259814eb4d687b6970"
12+
},
813
{
914
"name": "workspace",
1015
"source": "workspace"

dockerfile-optimizer/go.mod

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ go 1.23.2
55
require (
66
github.com/99designs/gqlgen v0.17.66
77
github.com/Khan/genqlient v0.8.0
8-
github.com/vektah/gqlparser/v2 v2.5.22
8+
github.com/vektah/gqlparser/v2 v2.5.23
99
go.opentelemetry.io/otel v1.32.0
1010
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0
1111
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.8.0
@@ -26,6 +26,7 @@ require (
2626

2727
require (
2828
github.com/cenkalti/backoff/v4 v4.3.0 // indirect
29+
github.com/dustin/go-humanize v1.0.1 // indirect
2930
github.com/go-logr/logr v1.4.2 // indirect
3031
github.com/go-logr/stdr v1.2.2 // indirect
3132
github.com/google/uuid v1.6.0 // indirect

dockerfile-optimizer/go.sum

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3
88
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
99
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
1010
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
11+
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
12+
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
1113
github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
1214
github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
1315
github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
@@ -29,8 +31,8 @@ github.com/sosodev/duration v1.3.1 h1:qtHBDMQ6lvMQsL15g4aopM4HEfOaYuhWBw3NPTtlqq
2931
github.com/sosodev/duration v1.3.1/go.mod h1:RQIBBX0+fMLc/D9+Jb/fwvVmo0eZvDDEERAikUR6SDg=
3032
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
3133
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
32-
github.com/vektah/gqlparser/v2 v2.5.22 h1:yaaeJ0fu+nv1vUMW0Hl+aS1eiv1vMfapBNjpffAda1I=
33-
github.com/vektah/gqlparser/v2 v2.5.22/go.mod h1:xMl+ta8a5M1Yo1A1Iwt/k7gSpscwSnHZdw7tfhEGfTM=
34+
github.com/vektah/gqlparser/v2 v2.5.23 h1:PurJ9wpgEVB7tty1seRUwkIDa/QH5RzkzraiKIjKLfA=
35+
github.com/vektah/gqlparser/v2 v2.5.23/go.mod h1:D1/VCZtV3LPnQrcPBeR/q5jkSQIPti0uYCP/RI0gIeo=
3436
go.opentelemetry.io/otel v1.32.0 h1:WnBN+Xjcteh0zdk01SVqV55d/m62NJLJdIyb4y/WO5U=
3537
go.opentelemetry.io/otel v1.32.0/go.mod h1:00DCVSB0RQcnzlwyTfqtxSm+DRr9hpYrHjNGiBHVQIg=
3638
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0 h1:WzNab7hOOLzdDF/EoWCt4glhrbMPVMOO5JYTmpz36Ls=

dockerfile-optimizer/main.go

Lines changed: 155 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,36 +16,174 @@ package main
1616

1717
import (
1818
"context"
19+
"encoding/json"
20+
"fmt"
21+
"path/filepath"
1922
"strings"
2023

2124
"dagger/dockerfile-optimizer/internal/dagger"
25+
26+
"github.com/dustin/go-humanize"
27+
"github.com/google/uuid"
2228
)
2329

2430
type DockerfileOptimizer struct{}
2531

32+
// Build the image from the Dockerfile, returns the number of layers and the size of the image
33+
func imageInfo(ctx context.Context, dir *dagger.Directory, path string) ([]int, error) {
34+
dirname, filename := filepath.Split(path)
35+
ctr := dag.Container().
36+
Build(dir.Directory(dirname), dagger.ContainerBuildOpts{Dockerfile: filename})
37+
38+
// Mount the OCI image and run tests
39+
out, err := dag.Container().From("wagoodman/dive:latest").
40+
WithMountedFile("/tmp/image.tar", ctr.AsTarball(
41+
// Layer compression seems to cause issues with dive in some cases
42+
dagger.ContainerAsTarballOpts{ForcedCompression: dagger.ImageLayerCompressionUncompressed},
43+
)).
44+
WithMountedDirectory("/workspace", dir.Directory(dirname)).
45+
WithExec([]string{"dive", "--json", "/tmp/img-info.json", "--ci", "docker-archive:///tmp/image.tar"}).
46+
File("/tmp/img-info.json").Contents(ctx)
47+
48+
if err != nil {
49+
return nil, fmt.Errorf("failed to run dive: %w", err)
50+
}
51+
52+
var imgInfo struct {
53+
Layer []struct{} `json:"layer"`
54+
Image struct {
55+
SizeBytes int64 `json:"sizeBytes"`
56+
} `json:"image"`
57+
}
58+
59+
if err := json.Unmarshal([]byte(out), &imgInfo); err != nil {
60+
return nil, fmt.Errorf("failed to parse image info: %w", err)
61+
}
62+
63+
numLayers := len(imgInfo.Layer)
64+
return []int{numLayers, int(imgInfo.Image.SizeBytes)}, nil
65+
}
66+
67+
func askLLM(ws *dagger.Workspace, dockerfile, extraContext string) *dagger.Llm {
68+
llm := dag.Llm().
69+
WithWorkspace(ws).
70+
WithPromptVar("dockerfile", dockerfile).
71+
WithPromptVar("extra_context", extraContext).
72+
WithPrompt(`
73+
You are a Platform Engineer with deep knowledge of Dockerfiles. You have access to a workspace.
74+
Use the read, write and build tools to complete the following assignment.
75+
76+
- Build the Dockerfile in the provided workspace at the path: "$dockerfile"
77+
- Optimize the Dockerfile for reducing its size, number of layers, and build time. And if possible, increasing the security level of the image by implementing best practices.
78+
- Ensure to not downgrade any version found in the Dockerfile.
79+
- If the Dockerfile is already optimized, just return an explanation that you couldn't optimize it.
80+
- Make sure the Dockerfile builds correctly before going to the next step.
81+
- If you have changes to make, write the optimized Dockerfile to the workspace, at the same path "$dockerfile".
82+
83+
At the end, return an explanation of the changes you made to the Dockerfile.
84+
$extra_context
85+
`)
86+
87+
return llm
88+
}
89+
90+
// Create a new PullRequest with the changes in the workspace, the given title and body, returns the PR URL
91+
func createPR(ctx context.Context, githubToken *dagger.Secret, repoURL string, w *dagger.Workspace, path, llmAnswer string) (string, error) {
92+
// generate a random branch name
93+
branchName := "dockerfile-improvements-" + uuid.New().String()[:8]
94+
// The changeset needs to contain only the Dockerfile otherwise the diff will fail (FIXME?)
95+
changeset := dag.Directory().WithFile(path, w.Workdir().File(path))
96+
// Create a new feature branch
97+
featureBranch := dag.FeatureBranch(githubToken, repoURL, branchName).
98+
WithChanges(changeset)
99+
100+
// Make sure changes have been made to the workspace
101+
diff, err := featureBranch.Diff(ctx)
102+
if err != nil {
103+
return "", err
104+
}
105+
106+
if diff == "" {
107+
return "", fmt.Errorf("got empty diff on feature branch (llm did not make any changes)")
108+
}
109+
110+
return featureBranch.PullRequest(ctx, "Optimizing Dockerfile", llmAnswer)
111+
}
112+
26113
// Optimize a Dockerfile
27114
func (m *DockerfileOptimizer) OptimizeDockerfile(ctx context.Context, githubToken *dagger.Secret, repoURL string) (string, error) {
28115
repoURL = strings.TrimSuffix(repoURL, ".git")
29116

30117
// Create a new workspace, using third-party module
31118
ws := dag.Workspace(githubToken, repoURL)
32-
// Run the agent loop in the workspace
33-
after := dag.Llm().
34-
WithWorkspace(ws).
35-
WithPrompt(`
36-
You are a Platform Engineer with deep knowledge of Dockerfiles. You have access to a workspace.
37-
Use the read, write, build, find, create-pr tools to complete the following assignment.
119+
originalWorkdir := ws.Workdir()
38120

39-
1. Look for a Dockerfile in the workspace (use the find tool with the "*Dockerfile*" pattern).
40-
2. Read the Dockerfile and optimize it for reducing its size, number of layers,
41-
and build time. And if possible, increasing the security level of the image by implementing best practices.
42-
3. Write the optimized Dockerfile to the workspace, in the same directory, replacing the original one.
43-
4. Build the container from the optimized Dockerfile to ensure it builds.
44-
5. Create a Pull Request with the changes in the workspace, with a useful title, the body should include the whole explanation of the changes made to the Dockerfile.
45-
6. Return the Pull Request URL as the last message.
121+
// Find the Dockerfile
122+
// FIXME: handle multiple Dockerfiles
123+
dockerfiles, err := ws.Workdir().Glob(ctx, "*Dockerfile*")
124+
if err != nil {
125+
return "", fmt.Errorf("cannot read the directory: %w", err)
126+
}
46127

47-
If the container build on step 4, read the build error and try to fix it in the Dockerfile until the container builds.
48-
`)
49-
// Return the last message which containers the Pull Request URL
50-
return after.LastReply(ctx)
128+
if len(dockerfiles) == 0 {
129+
return "", fmt.Errorf("no Dockerfile found")
130+
}
131+
132+
dockerfile := dockerfiles[0]
133+
134+
// Get the image info
135+
originalImgInfo, err := imageInfo(ctx, ws.Workdir(), dockerfile)
136+
if err != nil {
137+
return "", fmt.Errorf("failed to get image info: %w", err)
138+
}
139+
140+
extraContext := ""
141+
answer := ""
142+
var lastState *dagger.Workspace
143+
var lastImgInfo []int
144+
// Try 5 times to optimize the Dockerfile
145+
for range make([]int, 5) {
146+
// Ask the LLM to optimize the Dockerfile
147+
llm := askLLM(ws, dockerfile, extraContext)
148+
answer, err = llm.LastReply(ctx)
149+
if err != nil {
150+
return "", fmt.Errorf("failed to ask LLM: %w", err)
151+
}
152+
153+
lastState = llm.Workspace()
154+
155+
// Compare the optimized Dockerfile with the original one
156+
lastImgInfo, err = imageInfo(ctx, lastState.Workdir(), dockerfile)
157+
if err != nil {
158+
return "", fmt.Errorf("failed to get image info: %w", err)
159+
}
160+
161+
// We consider the optimization satisfactory if the size of the image is smaller
162+
if lastImgInfo[1] < originalImgInfo[1] {
163+
break
164+
}
165+
166+
// Otherwise we give extra context to the LLM and try again
167+
extraContext = "\n\nYou previously attempted to optimize the Dockerfile, but the changes were not satisfactory. Here are the details:\n\n"
168+
extraContext += fmt.Sprintf("- The number of layers is %d in the original image, and %d layers in the optimized version.\n", originalImgInfo[0], lastImgInfo[0])
169+
extraContext += fmt.Sprintf("- The original image size is %d bytes, and the optimized image size is %d bytes.\n\n", originalImgInfo[1], lastImgInfo[1])
170+
extraContext += "Please make the necessary changes to the Dockerfile to improve the image size and number of layers.\n"
171+
// FIXME: add the modified Dockerfile to the extra context?
172+
}
173+
174+
// Check if the workspace has been modified
175+
diff, err := originalWorkdir.Diff(lastState.Workdir()).Entries(ctx)
176+
if err != nil {
177+
return "", fmt.Errorf("failed to get workspace diff: %w", err)
178+
}
179+
180+
if len(diff) == 0 {
181+
return answer, fmt.Errorf("failed to optimize the Dockerfile")
182+
}
183+
184+
answer += "\n\nImage info:\n"
185+
answer += fmt.Sprintf("- The original image has %d layers and is %s in size.\n", originalImgInfo[0], humanize.Bytes(uint64(originalImgInfo[1])))
186+
answer += fmt.Sprintf("- The optimized image has %d layers and is %s in size.\n", lastImgInfo[0], humanize.Bytes(uint64(lastImgInfo[1])))
187+
188+
return createPR(ctx, githubToken, repoURL, lastState, dockerfile, answer)
51189
}
Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,7 @@
11
{
22
"name": "workspace",
3-
"engineVersion": "v0.17.0-llm.2",
3+
"engineVersion": "v0.17.0-llm.8",
44
"sdk": {
55
"source": "go"
6-
},
7-
"dependencies": [
8-
{
9-
"name": "feature-branch",
10-
"source": "github.com/kpenfound/dag/feature-branch",
11-
"pin": "c6809821f7133cc03bea4290fbe81212e25338c8"
12-
}
13-
]
6+
}
147
}

dockerfile-optimizer/workspace/go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ go 1.23.2
55
require (
66
github.com/99designs/gqlgen v0.17.66
77
github.com/Khan/genqlient v0.8.0
8-
github.com/vektah/gqlparser/v2 v2.5.22
8+
github.com/vektah/gqlparser/v2 v2.5.23
99
go.opentelemetry.io/otel v1.32.0
1010
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0
1111
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp v0.8.0

dockerfile-optimizer/workspace/go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ github.com/sosodev/duration v1.3.1 h1:qtHBDMQ6lvMQsL15g4aopM4HEfOaYuhWBw3NPTtlqq
2929
github.com/sosodev/duration v1.3.1/go.mod h1:RQIBBX0+fMLc/D9+Jb/fwvVmo0eZvDDEERAikUR6SDg=
3030
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
3131
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
32-
github.com/vektah/gqlparser/v2 v2.5.22 h1:yaaeJ0fu+nv1vUMW0Hl+aS1eiv1vMfapBNjpffAda1I=
33-
github.com/vektah/gqlparser/v2 v2.5.22/go.mod h1:xMl+ta8a5M1Yo1A1Iwt/k7gSpscwSnHZdw7tfhEGfTM=
32+
github.com/vektah/gqlparser/v2 v2.5.23 h1:PurJ9wpgEVB7tty1seRUwkIDa/QH5RzkzraiKIjKLfA=
33+
github.com/vektah/gqlparser/v2 v2.5.23/go.mod h1:D1/VCZtV3LPnQrcPBeR/q5jkSQIPti0uYCP/RI0gIeo=
3434
go.opentelemetry.io/otel v1.32.0 h1:WnBN+Xjcteh0zdk01SVqV55d/m62NJLJdIyb4y/WO5U=
3535
go.opentelemetry.io/otel v1.32.0/go.mod h1:00DCVSB0RQcnzlwyTfqtxSm+DRr9hpYrHjNGiBHVQIg=
3636
go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploggrpc v0.8.0 h1:WzNab7hOOLzdDF/EoWCt4glhrbMPVMOO5JYTmpz36Ls=

dockerfile-optimizer/workspace/main.go

Lines changed: 8 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,16 @@ package main
1616

1717
import (
1818
"context"
19-
"fmt"
2019
"path/filepath"
2120
"strings"
2221

2322
"dagger/workspace/internal/dagger"
24-
25-
"github.com/google/uuid"
2623
)
2724

2825
type Workspace struct {
29-
// The workspace's container state
26+
// The workspace's directory state
3027
// +internal-use-only
31-
Container *dagger.Container
28+
Workdir *dagger.Directory
3229
// Repository URL
3330
// +internal-use-only
3431
RepoURL string
@@ -47,75 +44,29 @@ func New(githubToken *dagger.Secret, repoURL string) Workspace {
4744

4845
return Workspace{
4946
// Build a base container optimized for Go development
50-
Container: dag.Container().
51-
From("cgr.dev/chainguard/wolfi-base").
52-
WithMountedDirectory("/src", repo).
53-
WithWorkdir("/src"),
47+
Workdir: repo,
5448
RepoURL: repoURL,
5549
GitHubToken: githubToken,
5650
}
5751
}
5852

59-
// func (w *Workspace) Ctr(ctx context.Context) *dagger.Container {
60-
// return w.Container
61-
// }
62-
63-
// Ensure the path is relative to the git clone directory
64-
func translatePath(path string) string {
65-
if strings.HasPrefix(path, "/src") {
66-
return path
67-
}
68-
return filepath.Join("/src", path)
69-
}
70-
7153
// Read a file at the given path
7254
func (w *Workspace) Read(ctx context.Context, path string) (string, error) {
73-
path = translatePath(path)
74-
return w.Container.File(path).Contents(ctx)
55+
return w.Workdir.File(path).Contents(ctx)
7556
}
7657

7758
// Write a file at the given path with the given content
7859
func (w Workspace) Write(path, content string) Workspace {
79-
path = translatePath(path)
80-
w.Container = w.Container.WithNewFile(path, content)
60+
w.Workdir = w.Workdir.WithNewFile(path, content)
8161
return w
8262
}
8363

8464
// Build the container from the Dockerfile at the given path
8565
func (w *Workspace) Build(ctx context.Context, path string) error {
8666
// Split directory and filename from path
8767
dirname, filename := filepath.Split(path)
88-
path = translatePath(path)
89-
_, err := w.Container.Build(w.Container.Directory(dirname), dagger.ContainerBuildOpts{Dockerfile: filename}).Sync(ctx)
68+
_, err := dag.Container().
69+
Build(w.Workdir.Directory(dirname), dagger.ContainerBuildOpts{Dockerfile: filename}).
70+
Sync(ctx)
9071
return err
9172
}
92-
93-
// Find files that match the given glob pattern
94-
func (w *Workspace) Find(ctx context.Context,
95-
pattern string,
96-
) ([]string, error) {
97-
return w.Container.Directory("/src").Glob(ctx, pattern)
98-
}
99-
100-
// Create a new PullRequest with the changes in the workspace, the given title and body, returns the PR URL
101-
func (w *Workspace) CreatePR(ctx context.Context, title, body string) (string, error) {
102-
// generate a random branch name
103-
branchName := "dockerfile-improvements-" + uuid.New().String()[:8]
104-
// The changeset needs to contain only the Dockerfile otherwise the diff will fail (FIXME?)
105-
changeset := dag.Directory().WithFile("Dockerfile", w.Container.File("Dockerfile"))
106-
// Create a new feature branch
107-
featureBranch := dag.FeatureBranch(w.GitHubToken, w.RepoURL, branchName).
108-
WithChanges(changeset)
109-
110-
// Make sure changes have been made to the workspace
111-
diff, err := featureBranch.Diff(ctx)
112-
if err != nil {
113-
return "", err
114-
}
115-
116-
if diff == "" {
117-
return "", fmt.Errorf("got empty diff on feature branch (llm did not make any changes)")
118-
}
119-
120-
return featureBranch.PullRequest(ctx, title, body)
121-
}

0 commit comments

Comments
 (0)