├── examples ├── hello-world ├── file-counter ├── backup-files ├── convert-audio ├── cleanup-old ├── process-logs └── product-downloader ├── go.mod ├── .gitignore ├── Makefile ├── .github └── workflows │ ├── release.yml │ ├── docker.yml │ └── ci.yml ├── internal ├── llm │ ├── script.go │ ├── types.go │ ├── base.go │ ├── prompts.go │ ├── provider.go │ ├── ollama.go │ ├── claude.go │ └── openai.go ├── test │ ├── test.go │ └── test_test.go ├── script │ ├── cache.go │ ├── pipeline_test.go │ ├── shell.go │ ├── executor.go │ └── pipeline.go ├── progress │ └── spinner.go ├── log │ └── log.go └── config │ ├── config.go │ └── config_test.go ├── Dockerfile ├── LICENSE ├── go.sum ├── README.md └── cmd └── llmscript └── main.go /examples/hello-world: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Print "Hello, world!" 4 | 5 | -------------------------------------------------------------------------------- /examples/file-counter: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Count all files in the current directory and its subdirectories 4 | Group them by file extension 5 | Print a summary showing the count for each extension 6 | Sort the results by count in descending order -------------------------------------------------------------------------------- /examples/backup-files: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Create a backup directory named 'backups' with today's date 4 | For each file in the current directory: 5 | - Create a copy in the backup directory 6 | - Add timestamp to the filename 7 | - Skip directories 8 | - Skip the backup directory itself 9 | Print a summary of backed up files -------------------------------------------------------------------------------- /examples/convert-audio: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Create an output directory named 'converted' 4 | For each audio file in the current directory: 5 | - Convert to MP3 format using ffmpeg 6 | - Set bitrate to 192k 7 | - Skip if already MP3 8 | - Skip if not an audio file 9 | - Show progress during conversion 10 | Print a summary of converted files -------------------------------------------------------------------------------- /examples/cleanup-old: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Find all files in the current directory and subdirectories that are: 4 | - Older than 30 days 5 | - Larger than 100MB 6 | - Not in a .git directory 7 | Print a list of files that will be deleted 8 | Ask for confirmation before deleting 9 | Delete the files if confirmed 10 | Print a summary of deleted files -------------------------------------------------------------------------------- /examples/process-logs: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Find all .log files in the current directory 4 | For each log file: 5 | - Extract lines containing "ERROR" or "WARNING" 6 | - Add the filename as a prefix to each line 7 | - Save to a new file named 'errors_' 8 | - Count the number of errors and warnings 9 | Print a summary of errors and warnings found -------------------------------------------------------------------------------- /examples/product-downloader: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env llmscript 2 | 3 | Create output directory `products` 4 | Download JSON product list from https://dummyjson.com/products 5 | Use jq to extract ID and image URL using '.products[].id, .products[].images[0]' 6 | Limit to 10 products 7 | Download each image to `products/.jpg` 8 | Resize each image to 100x100 using ImageMagick 9 | 10 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/statico/llmscript 2 | 3 | go 1.23.0 4 | 5 | toolchain go1.24.1 6 | 7 | require ( 8 | golang.org/x/term v0.30.0 9 | gopkg.in/yaml.v3 v3.0.1 10 | ) 11 | 12 | require golang.org/x/sys v0.31.0 // indirect 13 | 14 | require ( 15 | github.com/davecgh/go-spew v1.1.1 // indirect 16 | github.com/pmezard/go-difflib v1.0.0 // indirect 17 | github.com/stretchr/testify v1.10.0 18 | ) 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Binaries 2 | bin/ 3 | *.exe 4 | *.exe~ 5 | *.dll 6 | *.so 7 | *.dylib 8 | 9 | # Test binary, built with `go test -c` 10 | *.test 11 | 12 | # Output of the go coverage tool, specifically when used with LiteIDE 13 | *.out 14 | 15 | # Dependency directories 16 | vendor/ 17 | 18 | # Go workspace file 19 | go.work 20 | 21 | # IDE specific files 22 | .idea/ 23 | .vscode/ 24 | *.swp 25 | *.swo 26 | 27 | # OS specific files 28 | .DS_Store 29 | Thumbs.db -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: build test clean lint example 2 | 3 | build: 4 | go build -o bin/llmscript cmd/llmscript/main.go 5 | 6 | test: 7 | go test ./... 8 | 9 | clean: 10 | rm -rf bin/ 11 | go clean 12 | 13 | lint: 14 | @if ! command -v golangci-lint &> /dev/null; then \ 15 | echo "golangci-lint not found. Please install it with: brew install golangci-lint"; \ 16 | exit 1; \ 17 | fi 18 | golangci-lint run 19 | 20 | example: 21 | make 22 | ./bin/llmscript --verbose --no-cache examples/hello-world 23 | 24 | .DEFAULT_GOAL := build -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | on: 4 | push: 5 | tags: 6 | - "v*" # Trigger on version tags 7 | 8 | jobs: 9 | release: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - name: Checkout 13 | uses: actions/checkout@v4 14 | with: 15 | fetch-depth: 0 # Get all tags and history 16 | 17 | - name: Create Release 18 | uses: softprops/action-gh-release@v1 19 | with: 20 | generate_release_notes: true # Automatically generate release notes from commits 21 | draft: false 22 | prerelease: false 23 | -------------------------------------------------------------------------------- /internal/llm/script.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | import ( 4 | "strings" 5 | ) 6 | 7 | // ExtractScriptContent extracts the content between ") 13 | 14 | // If no tags found, return the entire response 15 | if start == -1 || end == -1 { 16 | return strings.TrimSpace(response) 17 | } 18 | 19 | // Extract content between tags and trim whitespace 20 | content := response[start+8 : end] 21 | return strings.TrimSpace(content) 22 | } 23 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | # Build stage 2 | FROM golang:1.24.1-alpine AS builder 3 | 4 | # Install build dependencies 5 | RUN apk add --no-cache git make 6 | 7 | # Set working directory 8 | WORKDIR /app 9 | 10 | # Copy go mod and sum files 11 | COPY go.mod go.sum ./ 12 | 13 | # Download dependencies 14 | RUN go mod download 15 | 16 | # Copy source code 17 | COPY . . 18 | 19 | # Build the application 20 | RUN make build 21 | 22 | # Final stage 23 | FROM alpine:latest 24 | 25 | # Install runtime dependencies 26 | RUN apk add --no-cache ca-certificates curl bash 27 | 28 | # Set working directory 29 | WORKDIR /app 30 | 31 | # Copy binary from builder and install it in PATH 32 | COPY --from=builder /app/bin/llmscript /usr/local/bin/ 33 | 34 | # Copy examples directory 35 | COPY --from=builder /app/examples ./examples 36 | 37 | # Set the default command 38 | CMD ["llmscript"] -------------------------------------------------------------------------------- /internal/llm/types.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | // OllamaConfig represents configuration for the Ollama provider 4 | type OllamaConfig struct { 5 | Model string `yaml:"model"` 6 | Host string `yaml:"host"` 7 | } 8 | 9 | // ClaudeConfig represents configuration for the Claude provider 10 | type ClaudeConfig struct { 11 | APIKey string `yaml:"api_key"` 12 | Model string `yaml:"model"` 13 | } 14 | 15 | // OpenAIConfig represents configuration for the OpenAI provider 16 | type OpenAIConfig struct { 17 | APIKey string `yaml:"api_key"` 18 | Model string `yaml:"model"` 19 | } 20 | 21 | // ProviderConfig represents the configuration for any LLM provider 22 | type ProviderConfig struct { 23 | Provider string `yaml:"provider"` 24 | Ollama OllamaConfig `yaml:"ollama,omitempty"` 25 | Claude ClaudeConfig `yaml:"claude,omitempty"` 26 | OpenAI OpenAIConfig `yaml:"openai,omitempty"` 27 | } 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Ian Langworth 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 2 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 3 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 4 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 5 | github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= 6 | github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= 7 | golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= 8 | golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= 9 | golang.org/x/term v0.30.0 h1:PQ39fJZ+mfadBm0y5WlL4vlM7Sx1Hgf13sMIY2+QS9Y= 10 | golang.org/x/term v0.30.0/go.mod h1:NYYFdzHoI5wRh/h5tDMdMqCqPJZEuNqVR5xJLd/n67g= 11 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 12 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 13 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 14 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 15 | -------------------------------------------------------------------------------- /.github/workflows/docker.yml: -------------------------------------------------------------------------------- 1 | name: Docker Build and Publish 2 | 3 | on: 4 | release: 5 | types: [published] 6 | 7 | env: 8 | REGISTRY: ghcr.io 9 | IMAGE_NAME: ${{ github.repository }} 10 | 11 | jobs: 12 | build-and-push: 13 | runs-on: ubuntu-latest 14 | permissions: 15 | contents: read 16 | packages: write 17 | 18 | steps: 19 | - name: Checkout repository 20 | uses: actions/checkout@v4 21 | 22 | - name: Set up Docker Buildx 23 | uses: docker/setup-buildx-action@v3 24 | 25 | - name: Log in to the Container registry 26 | uses: docker/login-action@v3 27 | with: 28 | registry: ${{ env.REGISTRY }} 29 | username: ${{ github.actor }} 30 | password: ${{ secrets.GITHUB_TOKEN }} 31 | 32 | - name: Extract metadata (tags, labels) for Docker 33 | id: meta 34 | uses: docker/metadata-action@v5 35 | with: 36 | images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} 37 | tags: | 38 | type=semver,pattern={{version}} 39 | type=semver,pattern={{major}}.{{minor}} 40 | type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }} 41 | 42 | - name: Build and push Docker image 43 | uses: docker/build-push-action@v5 44 | with: 45 | context: . 46 | push: true 47 | tags: ${{ steps.meta.outputs.tags }} 48 | labels: ${{ steps.meta.outputs.labels }} 49 | -------------------------------------------------------------------------------- /internal/test/test.go: -------------------------------------------------------------------------------- 1 | package test 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os" 7 | "os/exec" 8 | "path/filepath" 9 | ) 10 | 11 | // Test represents a test case for a script 12 | type Test struct { 13 | Name string 14 | Input string 15 | WantOutput string 16 | WantError bool 17 | WantExitErr bool 18 | } 19 | 20 | // TestRunner handles executing tests in a controlled environment 21 | type TestRunner struct { 22 | workDir string 23 | } 24 | 25 | // NewTestRunner creates a new test runner with a temporary working directory 26 | func NewTestRunner() (*TestRunner, error) { 27 | workDir, err := os.MkdirTemp("", "llmscript-test-*") 28 | if err != nil { 29 | return nil, fmt.Errorf("failed to create temp directory: %w", err) 30 | } 31 | return &TestRunner{workDir: workDir}, nil 32 | } 33 | 34 | // RunTest executes a test script in a controlled environment 35 | func (r *TestRunner) RunTest(ctx context.Context, mainScript, testScript string) error { 36 | // Create test directory 37 | testDir := filepath.Join(r.workDir, "test") 38 | if err := os.MkdirAll(testDir, 0755); err != nil { 39 | return fmt.Errorf("failed to create test directory: %w", err) 40 | } 41 | 42 | // Write scripts to files 43 | mainScriptPath := filepath.Join(testDir, "script.sh") 44 | testScriptPath := filepath.Join(testDir, "test.sh") 45 | 46 | if err := os.WriteFile(mainScriptPath, []byte(mainScript), 0750); err != nil { 47 | return fmt.Errorf("failed to write main script: %w", err) 48 | } 49 | if err := os.WriteFile(testScriptPath, []byte(testScript), 0750); err != nil { 50 | return fmt.Errorf("failed to write test script: %w", err) 51 | } 52 | 53 | // Run the test script 54 | cmd := exec.CommandContext(ctx, testScriptPath) 55 | cmd.Dir = testDir 56 | 57 | output, err := cmd.CombinedOutput() 58 | if err != nil { 59 | return fmt.Errorf("test failed: %w\nOutput:\n%s", err, output) 60 | } 61 | 62 | return nil 63 | } 64 | 65 | // Cleanup removes the temporary working directory 66 | func (r *TestRunner) Cleanup() error { 67 | return os.RemoveAll(r.workDir) 68 | } 69 | -------------------------------------------------------------------------------- /internal/test/test_test.go: -------------------------------------------------------------------------------- 1 | package test 2 | 3 | import ( 4 | "context" 5 | "os" 6 | "path/filepath" 7 | "testing" 8 | ) 9 | 10 | func TestTestRunner(t *testing.T) { 11 | runner, err := NewTestRunner() 12 | if err != nil { 13 | t.Fatalf("Failed to create test runner: %v", err) 14 | } 15 | defer func() { 16 | if err := runner.Cleanup(); err != nil { 17 | t.Errorf("failed to cleanup runner: %v", err) 18 | } 19 | }() 20 | 21 | ctx := context.Background() 22 | 23 | tests := []struct { 24 | name string 25 | mainScript string 26 | testScript string 27 | wantErr bool 28 | }{ 29 | { 30 | name: "simple echo test", 31 | mainScript: `#!/bin/bash 32 | echo "hello"`, 33 | testScript: `#!/bin/bash 34 | set -e 35 | [ "$(./script.sh)" = "hello" ] || exit 1`, 36 | wantErr: false, 37 | }, 38 | { 39 | name: "failing test", 40 | mainScript: `#!/bin/bash 41 | echo "wrong"`, 42 | testScript: `#!/bin/bash 43 | set -e 44 | [ "$(./script.sh)" = "right" ] || exit 1`, 45 | wantErr: true, 46 | }, 47 | { 48 | name: "test with setup", 49 | mainScript: `#!/bin/bash 50 | cat input.txt`, 51 | testScript: `#!/bin/bash 52 | set -e 53 | echo "test data" > input.txt 54 | [ "$(./script.sh)" = "test data" ] || exit 1`, 55 | wantErr: false, 56 | }, 57 | } 58 | 59 | for _, tt := range tests { 60 | t.Run(tt.name, func(t *testing.T) { 61 | err := runner.RunTest(ctx, tt.mainScript, tt.testScript) 62 | if (err != nil) != tt.wantErr { 63 | t.Errorf("RunTest() error = %v, wantErr %v", err, tt.wantErr) 64 | } 65 | }) 66 | } 67 | } 68 | 69 | func TestTestRunnerCleanup(t *testing.T) { 70 | runner, err := NewTestRunner() 71 | if err != nil { 72 | t.Fatalf("Failed to create test runner: %v", err) 73 | } 74 | 75 | // Create a file in the work directory 76 | testFile := filepath.Join(runner.workDir, "test.txt") 77 | if err := os.WriteFile(testFile, []byte("test"), 0644); err != nil { 78 | t.Fatalf("Failed to create test file: %v", err) 79 | } 80 | 81 | // Cleanup should remove the work directory 82 | if err := runner.Cleanup(); err != nil { 83 | t.Fatalf("Cleanup failed: %v", err) 84 | } 85 | 86 | // Verify the directory is gone 87 | if _, err := os.Stat(runner.workDir); !os.IsNotExist(err) { 88 | t.Error("Work directory still exists after cleanup") 89 | } 90 | } 91 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | tags: ["v*"] 7 | pull_request: 8 | branches: [main] 9 | 10 | jobs: 11 | test: 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: actions/checkout@v4 15 | 16 | - name: Set up Go 17 | uses: actions/setup-go@v5 18 | with: 19 | go-version: "1.24.1" 20 | check-latest: true 21 | 22 | - name: Install golangci-lint 23 | run: curl -sSfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(go env GOPATH)/bin v2.0.2 24 | 25 | - name: Add golangci-lint to PATH 26 | run: echo "$(go env GOPATH)/bin" >> $GITHUB_PATH 27 | 28 | - name: Install dependencies 29 | run: go mod download 30 | 31 | - name: Run tests 32 | run: go test ./... 33 | 34 | - name: Run linter 35 | run: golangci-lint run 36 | 37 | - name: Build 38 | run: make build 39 | 40 | release: 41 | needs: test 42 | runs-on: ubuntu-latest 43 | if: startsWith(github.ref, 'refs/tags/') 44 | 45 | permissions: 46 | contents: read 47 | packages: write 48 | 49 | steps: 50 | - uses: actions/checkout@v4 51 | with: 52 | fetch-depth: 0 53 | 54 | - name: Set up Docker Buildx 55 | uses: docker/setup-buildx-action@v3 56 | 57 | - name: Login to GitHub Container Registry 58 | uses: docker/login-action@v3 59 | with: 60 | registry: ghcr.io 61 | username: ${{ github.actor }} 62 | password: ${{ secrets.GITHUB_TOKEN }} 63 | 64 | - name: Extract metadata for Docker 65 | id: meta 66 | uses: docker/metadata-action@v5 67 | with: 68 | images: ghcr.io/${{ github.repository }} 69 | tags: | 70 | type=semver,pattern={{version}} 71 | type=semver,pattern={{major}}.{{minor}} 72 | type=semver,pattern={{major}} 73 | type=raw,value=latest,enable=${{ startsWith(github.ref, 'refs/tags/v') }} 74 | 75 | - name: Build and push Docker image 76 | uses: docker/build-push-action@v5 77 | with: 78 | context: . 79 | push: true 80 | tags: ${{ steps.meta.outputs.tags }} 81 | labels: ${{ steps.meta.outputs.labels }} 82 | -------------------------------------------------------------------------------- /internal/llm/base.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | ) 7 | 8 | // ScriptPair represents a feature script and its test script 9 | type ScriptPair struct { 10 | MainScript string // The feature script that implements the functionality 11 | TestScript string // The test script that verifies the feature script 12 | } 13 | 14 | // BaseProvider provides common functionality for LLM providers 15 | type BaseProvider struct { 16 | Config interface{} 17 | } 18 | 19 | // GenerateScripts generates a main script and test script pair 20 | func (p *BaseProvider) GenerateScripts(ctx context.Context, description string) (ScriptPair, error) { 21 | return ScriptPair{}, fmt.Errorf("GenerateScripts not implemented") 22 | } 23 | 24 | // FixScripts attempts to fix a script pair based on test failures 25 | func (p *BaseProvider) FixScripts(ctx context.Context, scripts ScriptPair, error string) (ScriptPair, error) { 26 | return ScriptPair{}, fmt.Errorf("FixScripts not implemented") 27 | } 28 | 29 | // GenerateScript is a default implementation that returns an error 30 | func (p *BaseProvider) GenerateScript(ctx context.Context, description string) (string, error) { 31 | return "", fmt.Errorf("GenerateScript not implemented") 32 | } 33 | 34 | // GenerateTests is a default implementation that returns an error 35 | func (p *BaseProvider) GenerateTests(ctx context.Context, script string, description string) ([]Test, error) { 36 | return nil, fmt.Errorf("GenerateTests not implemented") 37 | } 38 | 39 | // FixScript is a default implementation that returns an error 40 | func (p *BaseProvider) FixScript(ctx context.Context, script string, failures []TestFailure) (string, error) { 41 | return "", fmt.Errorf("FixScript not implemented") 42 | } 43 | 44 | // ValidateConfig validates the provider configuration 45 | func (p *BaseProvider) ValidateConfig() error { 46 | if p.Config == nil { 47 | return fmt.Errorf("config is required") 48 | } 49 | return nil 50 | } 51 | 52 | // formatPrompt formats a prompt template with the given arguments and platform information 53 | func (p *BaseProvider) formatPrompt(template string, args ...interface{}) string { 54 | // Add platform information as the last argument 55 | args = append(args, GetPlatformInfo()) 56 | return fmt.Sprintf(template, args...) 57 | } 58 | 59 | // Name returns a human-readable name for the provider 60 | func (p *BaseProvider) Name() string { 61 | return "" 62 | } 63 | -------------------------------------------------------------------------------- /internal/script/cache.go: -------------------------------------------------------------------------------- 1 | package script 2 | 3 | import ( 4 | "crypto/sha256" 5 | "encoding/hex" 6 | "encoding/json" 7 | "fmt" 8 | "os" 9 | "path/filepath" 10 | "strings" 11 | 12 | "github.com/statico/llmscript/internal/llm" 13 | ) 14 | 15 | // Cache handles caching of successful scripts and their test plans 16 | type Cache struct { 17 | dir string 18 | } 19 | 20 | // NewCache creates a new cache instance 21 | func NewCache() (*Cache, error) { 22 | homeDir, err := os.UserHomeDir() 23 | if err != nil { 24 | return nil, fmt.Errorf("failed to get home directory: %w", err) 25 | } 26 | 27 | cacheDir := filepath.Join(homeDir, ".config", "llmscript", "cache") 28 | if err := os.MkdirAll(cacheDir, 0755); err != nil { 29 | return nil, fmt.Errorf("failed to create cache directory: %w", err) 30 | } 31 | 32 | return &Cache{dir: cacheDir}, nil 33 | } 34 | 35 | // Get retrieves a cached script pair 36 | func (c *Cache) Get(description string) (llm.ScriptPair, error) { 37 | hash := c.hashDescription(description) 38 | scriptPath := filepath.Join(c.dir, hash+".json") 39 | 40 | // Check if file exists 41 | if _, err := os.Stat(scriptPath); os.IsNotExist(err) { 42 | return llm.ScriptPair{}, nil 43 | } 44 | 45 | // Read script pair 46 | data, err := os.ReadFile(scriptPath) 47 | if err != nil { 48 | return llm.ScriptPair{}, fmt.Errorf("failed to read cached script: %w", err) 49 | } 50 | 51 | var scripts llm.ScriptPair 52 | if err := json.Unmarshal(data, &scripts); err != nil { 53 | return llm.ScriptPair{}, fmt.Errorf("failed to parse cached scripts: %w", err) 54 | } 55 | 56 | return scripts, nil 57 | } 58 | 59 | // Set stores a successful script pair 60 | func (c *Cache) Set(description string, scripts llm.ScriptPair) error { 61 | hash := c.hashDescription(description) 62 | scriptPath := filepath.Join(c.dir, hash+".json") 63 | 64 | // Write script pair 65 | data, err := json.Marshal(scripts) 66 | if err != nil { 67 | return fmt.Errorf("failed to marshal scripts: %w", err) 68 | } 69 | if err := os.WriteFile(scriptPath, data, 0644); err != nil { 70 | return fmt.Errorf("failed to write cached script: %w", err) 71 | } 72 | 73 | return nil 74 | } 75 | 76 | // hashDescription generates a SHA-256 hash of the script description 77 | func (c *Cache) hashDescription(description string) string { 78 | hash := sha256.Sum256([]byte(strings.TrimSpace(description))) 79 | return hex.EncodeToString(hash[:]) 80 | } 81 | -------------------------------------------------------------------------------- /internal/progress/spinner.go: -------------------------------------------------------------------------------- 1 | package progress 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "os/signal" 7 | "strings" 8 | "sync/atomic" 9 | "syscall" 10 | "time" 11 | ) 12 | 13 | type Spinner struct { 14 | message atomic.Value 15 | messageWidth int 16 | parts []string 17 | value int 18 | ticker *time.Ticker 19 | started time.Time 20 | stopped time.Time 21 | sigChan chan os.Signal 22 | } 23 | 24 | func NewSpinner(message string) *Spinner { 25 | s := &Spinner{ 26 | parts: []string{ 27 | "⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏", 28 | }, 29 | started: time.Now(), 30 | sigChan: make(chan os.Signal, 1), 31 | } 32 | s.SetMessage(message) 33 | 34 | // Set up signal handling 35 | signal.Notify(s.sigChan, syscall.SIGINT, syscall.SIGTERM) 36 | 37 | go s.start() 38 | return s 39 | } 40 | 41 | func (s *Spinner) SetMessage(message string) { 42 | s.message.Store(message) 43 | } 44 | 45 | func (s *Spinner) String() string { 46 | var sb strings.Builder 47 | if s.stopped.IsZero() { 48 | spinner := s.parts[s.value] 49 | sb.WriteString(spinner) 50 | sb.WriteString(" ") 51 | } 52 | if message, ok := s.message.Load().(string); ok && len(message) > 0 { 53 | message := strings.TrimSpace(message) 54 | if s.messageWidth > 0 && len(message) > s.messageWidth { 55 | message = message[:s.messageWidth] 56 | } 57 | fmt.Fprintf(&sb, "%s", message) 58 | if padding := s.messageWidth - sb.Len(); padding > 0 { 59 | sb.WriteString(strings.Repeat(" ", padding)) 60 | } 61 | } 62 | return sb.String() 63 | } 64 | 65 | func (s *Spinner) start() { 66 | s.ticker = time.NewTicker(100 * time.Millisecond) 67 | fmt.Print("\r") // Start at beginning of line 68 | 69 | for { 70 | select { 71 | case <-s.ticker.C: 72 | s.value = (s.value + 1) % len(s.parts) 73 | if !s.stopped.IsZero() { 74 | return 75 | } 76 | fmt.Print("\r\033[2K") // Clear entire line 77 | fmt.Print(s.String()) 78 | case sig := <-s.sigChan: 79 | s.Stop() 80 | s.Clear() 81 | // Reset signal handling and re-send the signal 82 | signal.Stop(s.sigChan) 83 | p, err := os.FindProcess(os.Getpid()) 84 | if err == nil { 85 | if err := p.Signal(sig); err != nil { 86 | // Log the error but continue with the shutdown 87 | fmt.Fprintf(os.Stderr, "Error sending signal: %v\n", err) 88 | } 89 | } 90 | return 91 | } 92 | } 93 | } 94 | 95 | func (s *Spinner) Stop() { 96 | if s.stopped.IsZero() { 97 | s.stopped = time.Now() 98 | if s.ticker != nil { 99 | s.ticker.Stop() 100 | } 101 | } 102 | } 103 | 104 | // Clear clears the current line using ANSI escape sequences 105 | func (s *Spinner) Clear() { 106 | fmt.Print("\r\033[2K") // \033[2K clears the entire line 107 | } 108 | -------------------------------------------------------------------------------- /internal/script/pipeline_test.go: -------------------------------------------------------------------------------- 1 | package script 2 | 3 | import ( 4 | "context" 5 | "os" 6 | "path/filepath" 7 | "testing" 8 | "time" 9 | 10 | "github.com/statico/llmscript/internal/llm" 11 | "github.com/stretchr/testify/assert" 12 | "github.com/stretchr/testify/require" 13 | ) 14 | 15 | // mockLLMProvider implements the LLM provider interface for testing 16 | type mockLLMProvider struct { 17 | generateScriptsFunc func(ctx context.Context, description string) (llm.ScriptPair, error) 18 | fixScriptsFunc func(ctx context.Context, scripts llm.ScriptPair, error string) (llm.ScriptPair, error) 19 | } 20 | 21 | func (m *mockLLMProvider) GenerateScripts(ctx context.Context, description string) (llm.ScriptPair, error) { 22 | return m.generateScriptsFunc(ctx, description) 23 | } 24 | 25 | func (m *mockLLMProvider) FixScripts(ctx context.Context, scripts llm.ScriptPair, error string) (llm.ScriptPair, error) { 26 | return m.fixScriptsFunc(ctx, scripts, error) 27 | } 28 | 29 | func (m *mockLLMProvider) Name() string { 30 | return "mock" 31 | } 32 | 33 | func TestPipeline_GenerateAndTest(t *testing.T) { 34 | // Create a temporary directory for testing 35 | tmpDir, err := os.MkdirTemp("", "llmscript-test-*") 36 | if err != nil { 37 | t.Fatalf("failed to create temp dir: %v", err) 38 | } 39 | defer func() { 40 | if err := os.RemoveAll(tmpDir); err != nil { 41 | t.Errorf("failed to remove temp dir: %v", err) 42 | } 43 | }() 44 | 45 | // Create a mock LLM provider 46 | mockLLM := &mockLLMProvider{ 47 | generateScriptsFunc: func(ctx context.Context, description string) (llm.ScriptPair, error) { 48 | return llm.ScriptPair{ 49 | MainScript: `#!/bin/bash 50 | echo "Hello, World!"`, 51 | TestScript: `#!/bin/bash 52 | set -e 53 | [ "$(./script.sh)" = "Hello, World!" ] || exit 1`, 54 | }, nil 55 | }, 56 | fixScriptsFunc: func(ctx context.Context, scripts llm.ScriptPair, error string) (llm.ScriptPair, error) { 57 | return scripts, nil 58 | }, 59 | } 60 | 61 | // Create a new pipeline 62 | pipeline, err := NewPipeline(mockLLM, 1, 1, 5*time.Second, tmpDir, false) 63 | require.NoError(t, err) 64 | 65 | // Test script generation 66 | script, err := pipeline.GenerateAndTest(context.Background(), "Print 'Hello, World!'") 67 | require.NoError(t, err) 68 | 69 | // Verify the script was generated 70 | assert.Contains(t, script, "#!/bin/bash") 71 | assert.Contains(t, script, "echo \"Hello, World!\"") 72 | 73 | // Write the script to a file 74 | scriptPath := filepath.Join(tmpDir, "script.sh") 75 | err = os.WriteFile(scriptPath, []byte(script), 0755) 76 | require.NoError(t, err) 77 | 78 | // Verify the script file was created 79 | scriptContent, err := os.ReadFile(scriptPath) 80 | require.NoError(t, err) 81 | assert.Equal(t, script, string(scriptContent)) 82 | 83 | // Test script execution 84 | scripts := llm.ScriptPair{ 85 | MainScript: script, 86 | TestScript: `#!/bin/bash 87 | echo "Hello, World!"`, 88 | } 89 | err = pipeline.runTestScript(context.Background(), scripts) 90 | require.NoError(t, err) 91 | } 92 | -------------------------------------------------------------------------------- /internal/script/shell.go: -------------------------------------------------------------------------------- 1 | package script 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "os" 7 | "os/exec" 8 | "path/filepath" 9 | "runtime" 10 | "strings" 11 | ) 12 | 13 | // ShellConfig holds shell-specific configuration 14 | type ShellConfig struct { 15 | Path string 16 | Args []string 17 | IsLogin bool 18 | } 19 | 20 | // DetectShell determines the user's shell and returns its configuration 21 | func DetectShell() (*ShellConfig, error) { 22 | // Get the current user's shell from the environment or default 23 | shellPath := os.Getenv("SHELL") 24 | if shellPath == "" { 25 | if runtime.GOOS == "windows" { 26 | // On Windows, default to PowerShell if available, otherwise cmd.exe 27 | if path, err := exec.LookPath("powershell.exe"); err == nil { 28 | return &ShellConfig{ 29 | Path: path, 30 | Args: []string{"-NoProfile", "-NonInteractive", "-Command"}, 31 | }, nil 32 | } 33 | return &ShellConfig{ 34 | Path: "cmd.exe", 35 | Args: []string{"/C"}, 36 | }, nil 37 | } 38 | // On Unix-like systems, default to /bin/sh 39 | shellPath = "/bin/sh" 40 | } 41 | 42 | // Get the shell name from the path 43 | shellName := filepath.Base(shellPath) 44 | 45 | // Configure shell-specific arguments 46 | config := &ShellConfig{Path: shellPath} 47 | switch strings.ToLower(shellName) { 48 | case "bash": 49 | config.Args = []string{"--noprofile", "--norc", "-e", "-o", "pipefail", "-c"} 50 | case "zsh": 51 | config.Args = []string{"--no-rcs", "-e", "-o", "pipefail", "-c"} 52 | case "fish": 53 | config.Args = []string{"--no-config", "--command"} 54 | case "powershell.exe", "pwsh", "pwsh.exe": 55 | config.Args = []string{"-NoProfile", "-NonInteractive", "-Command"} 56 | default: 57 | // For unknown shells, use a simple -c argument 58 | config.Args = []string{"-c"} 59 | } 60 | 61 | return config, nil 62 | } 63 | 64 | // ValidateScript performs basic security checks on a script 65 | func ValidateScript(script string) error { 66 | // Check for suspicious commands or patterns 67 | suspiciousPatterns := []string{ 68 | "rm -rf /*", 69 | "mkfs", 70 | "> /dev/sd", 71 | "dd if=/dev/zero", 72 | ":(){:|:&};:", 73 | "wget", 74 | "curl.*| *sh", 75 | } 76 | 77 | for _, pattern := range suspiciousPatterns { 78 | if strings.Contains(strings.ToLower(script), strings.ToLower(pattern)) { 79 | return fmt.Errorf("script contains potentially dangerous pattern: %s", pattern) 80 | } 81 | } 82 | 83 | return nil 84 | } 85 | 86 | // PrepareScriptEnvironment sets up a secure environment for script execution 87 | func PrepareScriptEnvironment(workDir string) (string, error) { 88 | // Create a new temporary directory for script execution 89 | tmpDir, err := os.MkdirTemp(workDir, "script-") 90 | if err != nil { 91 | return "", fmt.Errorf("failed to create temporary directory: %w", err) 92 | } 93 | 94 | // Set restrictive permissions 95 | if err := os.Chmod(tmpDir, 0750); err != nil { 96 | if err := os.RemoveAll(tmpDir); err != nil { 97 | log.Printf("failed to remove temp dir: %v", err) 98 | } 99 | return "", fmt.Errorf("failed to set directory permissions: %w", err) 100 | } 101 | 102 | return tmpDir, nil 103 | } 104 | -------------------------------------------------------------------------------- /internal/log/log.go: -------------------------------------------------------------------------------- 1 | package log 2 | 3 | import ( 4 | "fmt" 5 | "os" 6 | "sync" 7 | 8 | "github.com/statico/llmscript/internal/progress" 9 | "golang.org/x/term" 10 | ) 11 | 12 | const ( 13 | red = "\033[31m" 14 | green = "\033[32m" 15 | yellow = "\033[33m" 16 | blue = "\033[34m" 17 | reset = "\033[0m" 18 | ) 19 | 20 | var ( 21 | level = InfoLevel 22 | levelLock sync.RWMutex 23 | spinner *progress.Spinner 24 | spinnerMu sync.Mutex 25 | ) 26 | 27 | type Level int 28 | 29 | const ( 30 | DebugLevel Level = iota 31 | InfoLevel 32 | WarnLevel 33 | ErrorLevel 34 | ) 35 | 36 | func SetLevel(l Level) { 37 | levelLock.Lock() 38 | defer levelLock.Unlock() 39 | level = l 40 | } 41 | 42 | // GetLevel returns the current logging level 43 | func GetLevel() Level { 44 | return getLevel() 45 | } 46 | 47 | func getLevel() Level { 48 | levelLock.RLock() 49 | defer levelLock.RUnlock() 50 | return level 51 | } 52 | 53 | func isTTY() bool { 54 | return term.IsTerminal(int(os.Stderr.Fd())) 55 | } 56 | 57 | // GetSpinner returns the singleton spinner instance, creating it if it doesn't exist 58 | func GetSpinner() *progress.Spinner { 59 | spinnerMu.Lock() 60 | defer spinnerMu.Unlock() 61 | if spinner == nil { 62 | spinner = progress.NewSpinner("") 63 | } 64 | return spinner 65 | } 66 | 67 | func updateSpinner(format string, args ...interface{}) { 68 | if !isTTY() { 69 | fmt.Fprintf(os.Stderr, format+"\n", args...) 70 | return 71 | } 72 | 73 | spinnerMu.Lock() 74 | defer spinnerMu.Unlock() 75 | 76 | if spinner == nil { 77 | spinner = progress.NewSpinner(fmt.Sprintf(format, args...)) 78 | } else { 79 | spinner.SetMessage(fmt.Sprintf(format, args...)) 80 | } 81 | fmt.Fprintf(os.Stderr, "\r%s", spinner) 82 | } 83 | 84 | func Spinner(format string, args ...interface{}) { 85 | updateSpinner(format, args...) 86 | } 87 | 88 | func Info(format string, args ...interface{}) { 89 | if getLevel() <= InfoLevel { 90 | if getLevel() == DebugLevel { 91 | fmt.Fprintf(os.Stderr, blue+"INFO: "+reset+format+"\n", args...) 92 | } else if isTTY() { 93 | updateSpinner(format, args...) 94 | } 95 | } 96 | } 97 | 98 | func Debug(format string, args ...interface{}) { 99 | if getLevel() <= DebugLevel { 100 | fmt.Fprintf(os.Stderr, blue+"DEBUG: "+reset+format+"\n", args...) 101 | } 102 | } 103 | 104 | func Warn(format string, args ...interface{}) { 105 | if getLevel() <= WarnLevel { 106 | fmt.Fprintf(os.Stderr, yellow+"WARN: "+reset+format+"\n", args...) 107 | } 108 | } 109 | 110 | func Error(format string, args ...interface{}) { 111 | if getLevel() <= ErrorLevel { 112 | fmt.Fprintf(os.Stderr, red+"ERROR: "+reset+format+"\n", args...) 113 | } 114 | } 115 | 116 | func Fatal(format string, args ...interface{}) { 117 | if spinner != nil { 118 | spinner.Stop() 119 | fmt.Fprintf(os.Stderr, "\n") 120 | } 121 | fmt.Fprintf(os.Stderr, red+"FATAL: "+reset+format+"\n", args...) 122 | os.Exit(1) 123 | } 124 | 125 | func Success(format string, args ...interface{}) { 126 | if spinner != nil { 127 | spinner.Stop() 128 | spinner.Clear() 129 | } 130 | if isTTY() || getLevel() == DebugLevel { 131 | fmt.Fprintf(os.Stderr, green+"✓ "+reset+format+"\n", args...) 132 | } 133 | } 134 | -------------------------------------------------------------------------------- /internal/llm/prompts.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | const ( 4 | // FeatureScriptPrompt is used to generate the main feature script 5 | FeatureScriptPrompt = `You are an expert shell script developer with deep knowledge of Unix/Linux systems, shell scripting best practices, and error handling. 6 | Your task is to create robust, maintainable shell scripts that work reliably across different environments. 7 | 8 | Create a shell script that accomplishes the following task: 9 | 10 | 11 | %s 12 | 13 | 14 | 15 | Target platform Information: 16 | %s 17 | 18 | 19 | 20 | - Use standard shell commands (sh/bash) with POSIX compliance where possible 21 | - Only use argument and environment variables if the description requires it 22 | - Follow shell scripting best practices 23 | - Ensure cross-platform compatibility and only use portable shell commands 24 | - Use proper exit codes for different scenarios 25 | - Keep the script short, concise, and simple 26 | 27 | 28 | 29 | Output your response in the following format: 30 | 31 | 35 | 36 | You *MUST NOT* include any other text, explanations, or markdown formatting. 37 | ` 38 | 39 | // TestScriptPrompt is used to generate the test script 40 | TestScriptPrompt = `You are an expert in testing shell scripts with extensive experience in test automation and quality assurance. 41 | Your goal is to create a test script that verifies the functionality of the main script, ./script.sh, including edge cases and error conditions. 42 | 43 | Create a test script for the following script: 44 | 45 | 48 | 49 | 50 | %s 51 | 52 | 53 | 54 | Target platform Information: 55 | %s 56 | 57 | 58 | 59 | - The script you're testing is ./script.sh 60 | - Create a test script that runs one or more test cases to make sure that ./script.sh works as expected 61 | - Each test case should: 62 | - Set up the test environment 63 | - Run the main script with test inputs 64 | - Verify the output matches expectations 65 | - Clean up after the test only if necessary 66 | - Do not modify ./script.sh, only test it 67 | - Do not use any randomization or nondeterministic functions 68 | - Use standard shell commands (sh/bash) with POSIX compliance where possible 69 | - The test script should not need any arguments to run 70 | - Return exit code 0 if all tests pass, or 1 if any test fails 71 | - Set appropriate timeouts for any long-running tests 72 | - Handle environment variables and cleanup 73 | - Ensure platform compatibility and only use portable shell commands 74 | - Keep the script short, concise, and simple 75 | 76 | 77 | 78 | Output your response in the following format: 79 | 80 | 84 | 85 | You *MUST NOT* include any other text, explanations, or markdown formatting. 86 | ` 87 | 88 | // FixScriptPrompt is used to fix a script based on test failures 89 | FixScriptPrompt = `You are an expert shell script developer specializing in debugging and fixing shell scripts. 90 | Your expertise includes error handling, cross-platform compatibility, and shell scripting best practices. 91 | 92 | Fix the following script based on the test failures: 93 | 94 | 97 | 98 | 99 | %s 100 | 101 | 102 | 103 | Target platform Information: 104 | %s 105 | 106 | 107 | 108 | - Fix all test failures while maintaining existing functionality 109 | - Improve error handling and validation 110 | - Follow shell scripting best practices 111 | - Ensure cross-platform compatibility 112 | - Keep the script short, concise, and simple 113 | 114 | 115 | 116 | Output your response in the following format: 117 | 118 | 122 | 123 | Do not include any other text, explanations, or markdown formatting. Only output the script between the markers. 124 | ` 125 | ) 126 | -------------------------------------------------------------------------------- /internal/llm/provider.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os/exec" 7 | "runtime" 8 | "strings" 9 | "time" 10 | ) 11 | 12 | // Test represents a test case for a script 13 | type Test struct { 14 | Name string 15 | Setup []string 16 | Input string 17 | Expected string 18 | Timeout time.Duration 19 | Environment map[string]string 20 | } 21 | 22 | // TestFailure represents a failed test case 23 | type TestFailure struct { 24 | Test Test 25 | Output string 26 | Error error 27 | ExitCode int 28 | } 29 | 30 | // Provider defines the interface for LLM providers 31 | type Provider interface { 32 | // GenerateScripts creates a main script and test script from a natural language description 33 | GenerateScripts(ctx context.Context, description string) (ScriptPair, error) 34 | // FixScripts attempts to fix both scripts based on test failures 35 | FixScripts(ctx context.Context, scripts ScriptPair, error string) (ScriptPair, error) 36 | // Name returns a human-readable name for the provider 37 | Name() string 38 | } 39 | 40 | // GetPlatformInfo returns information about the current platform 41 | func GetPlatformInfo() string { 42 | info := []string{ 43 | "Operating System: " + runtime.GOOS, 44 | "Architecture: " + runtime.GOARCH, 45 | } 46 | 47 | // Get additional system information using uname 48 | cmd := exec.Command("uname", "-a") 49 | if output, err := cmd.Output(); err == nil { 50 | info = append(info, "System Info: "+string(output)) 51 | } 52 | 53 | // Get shell information 54 | cmd = exec.Command("bash", "--version") 55 | if output, err := cmd.Output(); err == nil { 56 | info = append(info, "Shell Info: "+string(output)) 57 | } 58 | 59 | return strings.Join(info, "\n") 60 | } 61 | 62 | // NewProvider creates a new LLM provider based on the provider name 63 | func NewProvider(name string, config interface{}) (Provider, error) { 64 | if name == "" { 65 | name = "ollama" // Default to Ollama if no provider specified 66 | } 67 | 68 | // Default Ollama config 69 | ollamaConfig := OllamaConfig{ 70 | Model: "llama3.2", 71 | Host: "http://localhost:11434", 72 | } 73 | 74 | // If config is provided, try to extract values 75 | if config != nil { 76 | // Try to convert the config to a map[string]interface{} first 77 | if cfgMap, ok := config.(map[string]interface{}); ok { 78 | // Handle Ollama config 79 | if ollamaCfg, ok := cfgMap["ollama"].(map[string]interface{}); ok { 80 | if model, ok := ollamaCfg["model"].(string); ok && model != "" { 81 | ollamaConfig.Model = model 82 | } 83 | if host, ok := ollamaCfg["host"].(string); ok && host != "" { 84 | ollamaConfig.Host = host 85 | } 86 | } 87 | } 88 | } 89 | 90 | switch name { 91 | case "ollama": 92 | return NewOllamaProvider(ollamaConfig) 93 | case "claude": 94 | if config == nil { 95 | return nil, fmt.Errorf("a Claude API key is required") 96 | } 97 | // Try to convert the config to a map[string]interface{} 98 | cfgMap, ok := config.(map[string]interface{}) 99 | if !ok { 100 | return nil, fmt.Errorf("invalid config type for Claude provider") 101 | } 102 | // Extract Claude config 103 | claudeCfg, ok := cfgMap["claude"].(map[string]interface{}) 104 | if !ok { 105 | return nil, fmt.Errorf("missing Claude configuration") 106 | } 107 | apiKey, ok := claudeCfg["api_key"].(string) 108 | if !ok || apiKey == "" { 109 | return nil, fmt.Errorf("a Claude API key is required") 110 | } 111 | model, _ := claudeCfg["model"].(string) 112 | if model == "" { 113 | model = "claude-3-opus-20240229" // Default model 114 | } 115 | return NewClaudeProvider(ClaudeConfig{ 116 | APIKey: apiKey, 117 | Model: model, 118 | }) 119 | case "openai": 120 | if config == nil { 121 | return nil, fmt.Errorf("an OpenAI API key is required") 122 | } 123 | // Try to convert the config to a map[string]interface{} 124 | cfgMap, ok := config.(map[string]interface{}) 125 | if !ok { 126 | return nil, fmt.Errorf("invalid config type for OpenAI provider") 127 | } 128 | // Extract OpenAI config 129 | openaiCfg, ok := cfgMap["openai"].(map[string]interface{}) 130 | if !ok { 131 | return nil, fmt.Errorf("missing OpenAI configuration") 132 | } 133 | apiKey, ok := openaiCfg["api_key"].(string) 134 | if !ok || apiKey == "" { 135 | return nil, fmt.Errorf("an OpenAI API key is required") 136 | } 137 | model, _ := openaiCfg["model"].(string) 138 | if model == "" { 139 | model = "gpt-4-turbo-preview" // Default model 140 | } 141 | return NewOpenAIProvider(OpenAIConfig{ 142 | APIKey: apiKey, 143 | Model: model, 144 | }) 145 | default: 146 | return nil, fmt.Errorf("unsupported provider: %s", name) 147 | } 148 | } 149 | -------------------------------------------------------------------------------- /internal/llm/ollama.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "fmt" 7 | "io" 8 | "net/http" 9 | "strings" 10 | 11 | "github.com/statico/llmscript/internal/log" 12 | ) 13 | 14 | // OllamaProvider implements the Provider interface using Ollama 15 | type OllamaProvider struct { 16 | BaseProvider 17 | config OllamaConfig 18 | } 19 | 20 | // NewOllamaProvider creates a new Ollama provider 21 | func NewOllamaProvider(config OllamaConfig) (*OllamaProvider, error) { 22 | return &OllamaProvider{ 23 | BaseProvider: BaseProvider{Config: config}, 24 | config: config, 25 | }, nil 26 | } 27 | 28 | // GenerateScripts creates a main script and test script from a natural language description 29 | func (p *OllamaProvider) GenerateScripts(ctx context.Context, description string) (ScriptPair, error) { 30 | // First generate the main script 31 | log.Info("Generating main script with Ollama...") 32 | mainPrompt := p.formatPrompt(FeatureScriptPrompt, description) 33 | mainScript, err := p.generate(ctx, mainPrompt) 34 | if err != nil { 35 | return ScriptPair{}, fmt.Errorf("failed to generate main script: %w", err) 36 | } 37 | mainScript = ExtractScriptContent(mainScript) 38 | log.Debug("Main script generated:\n%s", mainScript) 39 | 40 | // Then generate the test script 41 | log.Info("Generating test script with Ollama...") 42 | testPrompt := p.formatPrompt(TestScriptPrompt, mainScript, description) 43 | testScript, err := p.generate(ctx, testPrompt) 44 | if err != nil { 45 | return ScriptPair{}, fmt.Errorf("failed to generate test script: %w", err) 46 | } 47 | testScript = ExtractScriptContent(testScript) 48 | log.Debug("Test script generated:\n%s", testScript) 49 | 50 | return ScriptPair{ 51 | MainScript: strings.TrimSpace(mainScript), 52 | TestScript: strings.TrimSpace(testScript), 53 | }, nil 54 | } 55 | 56 | // FixScripts attempts to fix both scripts based on test failures 57 | func (p *OllamaProvider) FixScripts(ctx context.Context, scripts ScriptPair, error string) (ScriptPair, error) { 58 | // Only fix the main script 59 | mainPrompt := p.formatPrompt(FixScriptPrompt, scripts.MainScript, error) 60 | fixedMainScript, err := p.generate(ctx, mainPrompt) 61 | if err != nil { 62 | return ScriptPair{}, fmt.Errorf("failed to fix main script: %w", err) 63 | } 64 | fixedMainScript = ExtractScriptContent(fixedMainScript) 65 | 66 | return ScriptPair{ 67 | MainScript: strings.TrimSpace(fixedMainScript), 68 | TestScript: scripts.TestScript, 69 | }, nil 70 | } 71 | 72 | // generate sends a prompt to Ollama and returns the response 73 | func (p *OllamaProvider) generate(ctx context.Context, prompt string) (string, error) { 74 | url := fmt.Sprintf("%s/api/generate", p.config.Host) 75 | 76 | reqBody := map[string]interface{}{ 77 | "model": p.config.Model, 78 | "prompt": prompt, 79 | "stream": false, 80 | } 81 | 82 | jsonData, err := json.Marshal(reqBody) 83 | if err != nil { 84 | return "", fmt.Errorf("failed to marshal request: %w", err) 85 | } 86 | 87 | req, err := http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(jsonData))) 88 | if err != nil { 89 | return "", fmt.Errorf("failed to create request: %w", err) 90 | } 91 | req.Header.Set("Content-Type", "application/json") 92 | 93 | resp, err := http.DefaultClient.Do(req) 94 | if err != nil { 95 | return "", fmt.Errorf("failed to send request: %w", err) 96 | } 97 | defer func() { 98 | if err := resp.Body.Close(); err != nil { 99 | log.Error("failed to close response body: %v", err) 100 | } 101 | }() 102 | 103 | if resp.StatusCode != http.StatusOK { 104 | body, _ := io.ReadAll(resp.Body) 105 | return "", fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(body)) 106 | } 107 | 108 | // Read the entire response body 109 | body, err := io.ReadAll(resp.Body) 110 | if err != nil { 111 | return "", fmt.Errorf("failed to read response body: %w", err) 112 | } 113 | 114 | // Parse the response 115 | var result struct { 116 | Response string `json:"response"` 117 | Done bool `json:"done"` 118 | } 119 | if err := json.Unmarshal(body, &result); err != nil { 120 | return "", fmt.Errorf("failed to decode response: %w", err) 121 | } 122 | 123 | return result.Response, nil 124 | } 125 | 126 | // Name returns a human-readable name for the provider 127 | func (p *OllamaProvider) Name() string { 128 | return "Ollama" 129 | } 130 | 131 | // formatFailures formats test failures into a string 132 | func formatFailures(failures []TestFailure) string { 133 | var parts []string 134 | for _, f := range failures { 135 | parts = append(parts, fmt.Sprintf("Test: %s\nOutput: %s\nError: %v\nExit Code: %d", 136 | f.Test.Name, f.Output, f.Error, f.ExitCode)) 137 | } 138 | return strings.Join(parts, "\n\n") 139 | } 140 | -------------------------------------------------------------------------------- /internal/script/executor.go: -------------------------------------------------------------------------------- 1 | package script 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os" 7 | "os/exec" 8 | "path/filepath" 9 | "strings" 10 | 11 | "github.com/statico/llmscript/internal/llm" 12 | "github.com/statico/llmscript/internal/log" 13 | ) 14 | 15 | // Executor handles running scripts in a controlled environment 16 | type Executor struct { 17 | workDir string 18 | shell *ShellConfig 19 | } 20 | 21 | // NewExecutor creates a new script executor 22 | func NewExecutor(workDir string) (*Executor, error) { 23 | if err := os.MkdirAll(workDir, 0755); err != nil { 24 | return nil, fmt.Errorf("failed to create work directory: %w", err) 25 | } 26 | 27 | shell, err := DetectShell() 28 | if err != nil { 29 | return nil, fmt.Errorf("failed to detect shell: %w", err) 30 | } 31 | 32 | return &Executor{ 33 | workDir: workDir, 34 | shell: shell, 35 | }, nil 36 | } 37 | 38 | // ExecuteTest runs a single test case in a controlled environment 39 | func (e *Executor) ExecuteTest(ctx context.Context, script string, test llm.Test) (string, error) { 40 | // Validate script for security 41 | if err := ValidateScript(script); err != nil { 42 | return "", fmt.Errorf("script validation failed: %w", err) 43 | } 44 | 45 | // Create a secure temporary directory for this test 46 | testDir, err := os.MkdirTemp("", "llmscript-test-*") 47 | if err != nil { 48 | return "", fmt.Errorf("failed to create test directory: %w", err) 49 | } 50 | defer func() { 51 | if err := os.RemoveAll(testDir); err != nil { 52 | log.Error("failed to remove test directory: %v", err) 53 | } 54 | }() 55 | 56 | log.Debug("Test directory: %s", testDir) 57 | 58 | // Write the script to a file 59 | scriptPath := filepath.Join(testDir, "script.sh") 60 | if err := os.WriteFile(scriptPath, []byte(script), 0750); err != nil { 61 | return "", fmt.Errorf("failed to write script: %w", err) 62 | } 63 | log.Debug("Script written to: %s", scriptPath) 64 | 65 | // Run setup commands 66 | for _, cmd := range test.Setup { 67 | log.Debug("Running setup command: %s", cmd) 68 | if err := e.runCommand(ctx, testDir, cmd, test.Environment); err != nil { 69 | return "", fmt.Errorf("setup command failed: %w", err) 70 | } 71 | } 72 | 73 | // Run the script with the test input 74 | args := append(e.shell.Args, scriptPath) 75 | cmd := exec.CommandContext(ctx, e.shell.Path, args...) 76 | cmd.Dir = testDir 77 | cmd.Env = e.buildEnv(test.Environment) 78 | 79 | log.Debug("Running script with shell: %s %v", e.shell.Path, args) 80 | if len(test.Environment) > 0 { 81 | log.Debug("Environment variables:") 82 | for k, v := range test.Environment { 83 | log.Debug(" %s=%s", k, v) 84 | } 85 | } 86 | 87 | // Set up input/output pipes 88 | stdin, err := cmd.StdinPipe() 89 | if err != nil { 90 | return "", fmt.Errorf("failed to get stdin pipe: %w", err) 91 | } 92 | defer func() { 93 | if err := stdin.Close(); err != nil { 94 | log.Error("failed to close stdin: %v", err) 95 | } 96 | }() 97 | 98 | var stdout, stderr strings.Builder 99 | cmd.Stdout = &stdout 100 | cmd.Stderr = &stderr 101 | 102 | // Start the command 103 | if err := cmd.Start(); err != nil { 104 | return "", fmt.Errorf("failed to start command: %w", err) 105 | } 106 | 107 | // Write input and close stdin 108 | if _, err := stdin.Write([]byte(test.Input)); err != nil { 109 | return "", fmt.Errorf("failed to write input: %w", err) 110 | } 111 | if err := stdin.Close(); err != nil { 112 | return "", fmt.Errorf("failed to close stdin: %w", err) 113 | } 114 | 115 | // Wait for completion with timeout 116 | done := make(chan error) 117 | go func() { 118 | done <- cmd.Wait() 119 | }() 120 | 121 | select { 122 | case err := <-done: 123 | if err != nil { 124 | return "", fmt.Errorf("command failed: %w\nstdout: %s\nstderr: %s", err, stdout.String(), stderr.String()) 125 | } 126 | case <-ctx.Done(): 127 | if cmd.Process != nil { 128 | if err := cmd.Process.Kill(); err != nil { 129 | log.Error("failed to kill process: %v", err) 130 | } 131 | } 132 | return "", fmt.Errorf("command timed out") 133 | } 134 | 135 | return stdout.String(), nil 136 | } 137 | 138 | // runCommand executes a shell command in the given directory 139 | func (e *Executor) runCommand(ctx context.Context, dir, cmd string, env map[string]string) error { 140 | args := append(e.shell.Args, cmd) 141 | command := exec.CommandContext(ctx, e.shell.Path, args...) 142 | command.Dir = dir 143 | command.Env = e.buildEnv(env) 144 | 145 | output, err := command.CombinedOutput() 146 | if err != nil { 147 | return fmt.Errorf("command failed: %w\noutput: %s", err, output) 148 | } 149 | return nil 150 | } 151 | 152 | // buildEnv builds the environment variables for a command 153 | func (e *Executor) buildEnv(env map[string]string) []string { 154 | // Start with current environment 155 | result := os.Environ() 156 | 157 | // Add test-specific environment variables 158 | for k, v := range env { 159 | result = append(result, fmt.Sprintf("%s=%s", k, v)) 160 | } 161 | 162 | return result 163 | } 164 | -------------------------------------------------------------------------------- /internal/script/pipeline.go: -------------------------------------------------------------------------------- 1 | package script 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "os" 7 | "os/exec" 8 | "path/filepath" 9 | "time" 10 | 11 | "github.com/statico/llmscript/internal/llm" 12 | "github.com/statico/llmscript/internal/log" 13 | ) 14 | 15 | // Test represents a test case for a script 16 | type Test = llm.Test 17 | 18 | // TestFailure represents a failed test case 19 | type TestFailure = llm.TestFailure 20 | 21 | // Pipeline handles the script generation and testing process 22 | type Pipeline struct { 23 | llm llm.Provider 24 | maxFixes int 25 | maxAttempts int 26 | timeout time.Duration 27 | workDir string 28 | cache *Cache 29 | noCache bool 30 | } 31 | 32 | // NewPipeline creates a new script generation pipeline 33 | func NewPipeline(llm llm.Provider, maxFixes, maxAttempts int, timeout time.Duration, workDir string, noCache bool) (*Pipeline, error) { 34 | if err := os.MkdirAll(workDir, 0755); err != nil { 35 | return nil, fmt.Errorf("failed to create work directory: %w", err) 36 | } 37 | 38 | var cache *Cache 39 | if !noCache { 40 | var err error 41 | cache, err = NewCache() 42 | if err != nil { 43 | return nil, fmt.Errorf("failed to create cache: %w", err) 44 | } 45 | } 46 | 47 | return &Pipeline{ 48 | llm: llm, 49 | maxFixes: maxFixes, 50 | maxAttempts: maxAttempts, 51 | timeout: timeout, 52 | workDir: workDir, 53 | cache: cache, 54 | noCache: noCache, 55 | }, nil 56 | } 57 | 58 | // GenerateAndTest generates a script from a natural language description and tests it 59 | func (p *Pipeline) GenerateAndTest(ctx context.Context, description string) (string, error) { 60 | // Check cache first if enabled 61 | if !p.noCache && p.cache != nil { 62 | log.Info("Checking cache...") 63 | if scripts, err := p.cache.Get(description); err == nil && scripts.MainScript != "" { 64 | // Run test script to verify 65 | if err := p.runTestScript(ctx, scripts); err == nil { 66 | log.Success("Cached script found") 67 | return scripts.MainScript, nil 68 | } 69 | log.Warn("Cached scripts failed verification, generating new scripts") 70 | } 71 | } 72 | 73 | // Generate initial scripts 74 | log.Info("Generating initial scripts with %s...", p.llm.Name()) 75 | scripts, err := p.llm.GenerateScripts(ctx, description) 76 | if err != nil { 77 | return "", fmt.Errorf("failed to generate initial scripts: %w", err) 78 | } 79 | log.Debug("Initial scripts generated") 80 | 81 | // Run test script and fix failures 82 | for attempt := 0; attempt < p.maxAttempts; attempt++ { 83 | if attempt > 0 { 84 | log.Info("Attempt %d/%d: Generating new scripts...", attempt+1, p.maxAttempts) 85 | scripts, err = p.llm.GenerateScripts(ctx, description) 86 | if err != nil { 87 | return "", fmt.Errorf("failed to generate new scripts: %w", err) 88 | } 89 | log.Debug("New scripts generated") 90 | } 91 | 92 | // Try to fix any failures 93 | for fix := 0; fix < p.maxFixes; fix++ { 94 | // Run test script 95 | log.Info("Testing script (attempt %d/%d)...", attempt+1, p.maxAttempts) 96 | err := p.runTestScript(ctx, scripts) 97 | if err == nil { 98 | // Cache successful scripts if caching is enabled 99 | if !p.noCache && p.cache != nil { 100 | log.Info("Caching successful scripts...") 101 | if err := p.cache.Set(description, scripts); err != nil { 102 | log.Warn("Failed to cache successful scripts: %v", err) 103 | } 104 | } 105 | return scripts.MainScript, nil 106 | } 107 | 108 | if fix < p.maxFixes-1 { // Don't try to fix on the last iteration 109 | log.Info("Fix attempt %d/%d...", fix+1, p.maxFixes) 110 | scripts, err = p.llm.FixScripts(ctx, scripts, err.Error()) 111 | if err != nil { 112 | return "", fmt.Errorf("failed to fix scripts: %w", err) 113 | } 114 | log.Debug("Scripts fixed") 115 | log.Debug("New script:\n%s", scripts.MainScript) 116 | } 117 | } 118 | } 119 | 120 | return "", fmt.Errorf("failed to generate working scripts after %d attempts", p.maxAttempts) 121 | } 122 | 123 | // runTestScript executes the test script in a controlled environment 124 | func (p *Pipeline) runTestScript(ctx context.Context, scripts llm.ScriptPair) error { 125 | // Create a secure temporary directory for this test 126 | testDir, err := os.MkdirTemp("", "llmscript-test-*") 127 | if err != nil { 128 | return fmt.Errorf("failed to create test directory: %w", err) 129 | } 130 | defer func() { 131 | if err := os.RemoveAll(testDir); err != nil { 132 | log.Error("failed to remove test directory: %v", err) 133 | } 134 | }() 135 | 136 | // Write both scripts to files 137 | featureScriptPath := filepath.Join(testDir, "script.sh") 138 | testScriptPath := filepath.Join(testDir, "test.sh") 139 | 140 | if err := os.WriteFile(featureScriptPath, []byte(scripts.MainScript), 0750); err != nil { 141 | return fmt.Errorf("failed to write feature script: %w", err) 142 | } 143 | if err := os.WriteFile(testScriptPath, []byte(scripts.TestScript), 0750); err != nil { 144 | return fmt.Errorf("failed to write test script: %w", err) 145 | } 146 | 147 | // Run the test script with timeout 148 | ctx, cancel := context.WithTimeout(ctx, p.timeout) 149 | defer cancel() 150 | 151 | cmd := exec.CommandContext(ctx, testScriptPath) 152 | cmd.Dir = testDir 153 | 154 | output, err := cmd.CombinedOutput() 155 | log.Debug("Test script output:\n%s", output) 156 | if err != nil { 157 | if exitErr, ok := err.(*exec.ExitError); ok { 158 | log.Debug("Test script exited with code: %d", exitErr.ExitCode()) 159 | } 160 | return fmt.Errorf("test script failed: %w\nOutput:\n%s", err, output) 161 | } 162 | log.Debug("Test script exited with code: 0") 163 | 164 | return nil 165 | } 166 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # llmscript 2 | 3 | [![CI](https://github.com/statico/llmscript/actions/workflows/ci.yml/badge.svg)](https://github.com/statico/llmscript/actions/workflows/ci.yml) 4 | 5 | llmscript is a shell script that uses a large language model (LLM) to build and test shell programs so that you can write scripts in natural language instead of bash or other shell scripting languages. 6 | 7 | a terminal window showing a demonstration of the llmscript tool to print hello world 8 | 9 | You can configure it to use [Ollama](https://ollama.com/) (free and local), [Claude](https://www.anthropic.com/claude) (paid), or [OpenAI](https://openai.com/) (paid). 10 | 11 | > [!NOTE] 12 | > Does this actually work? Yeah, somewhat! Could it create scripts that erase your drive? Maybe! Good luck! 13 | > 14 | > Most of this project was written by [Claude](https://www.anthropic.com/claude) with [Cursor](https://www.cursor.com). I can't actually claim that I "wrote" any of the source code. I barely know Go. 15 | 16 | ## Example 17 | 18 | ``` 19 | #!/usr/bin/env llmscript 20 | 21 | Create an output directory, `output`. 22 | For every PNG file in `input`: 23 | - Convert it to 256x256 with ImageMagick 24 | - Run pngcrush 25 | ``` 26 | 27 | Running it with a directory of PNG files would look like this: 28 | 29 | ``` 30 | $ ./convert-pngs 31 | ✓ Script generated successfully! 32 | Creating output directory 33 | Convering input/1.png 34 | Convering input/2.png 35 | Convering input/3.png 36 | Running pngcrush on output/1.png 37 | Running pngcrush on output/2.png 38 | Running pngcrush on output/3.png 39 | Done! 40 | ``` 41 | 42 | Running it again will use the cache and not generate any new scripts: 43 | 44 | ``` 45 | $ ./convert-pngs 46 | ✓ Cached script found 47 | Creating output directory 48 | Convering input/1.png 49 | ... 50 | ``` 51 | 52 | If you want to generate a new script, use the `--no-cache` flag. 53 | 54 | ## Prerequisites 55 | 56 | - [Go](https://go.dev/) (1.22 or later) 57 | - One of: 58 | - [Ollama](https://ollama.com/) running locally 59 | - A [Claude](https://www.anthropic.com/claude) API key 60 | - An [OpenAI](https://openai.com/) API key 61 | 62 | ## Installation 63 | 64 | ``` 65 | go install github.com/statico/llmscript/cmd/llmscript@latest 66 | ``` 67 | 68 | (Can't find it? Check `~/go/bin`.) 69 | 70 | Or, if you're spooked by running LLM-generated shell scripts (good for you!), consider running llmscript via Docker: 71 | 72 | ``` 73 | docker run --network host -it -v "$(pwd):/data" -w /data ghcr.io/statico/llmscript --verbose examples/hello-world 74 | ``` 75 | 76 | ## Usage 77 | 78 | Create a script file like the above example, or check out the [examples](examples) directory for more. You can use a shebang like: 79 | 80 | ``` 81 | #!/usr/bin/env llmscript 82 | ``` 83 | 84 | or run it directly like: 85 | 86 | ``` 87 | $ llmscript hello-world 88 | ``` 89 | 90 | By default, llmscript will use Ollama with the `llama3.2` model. You can configure this by creating a config file with the `llmscript --write-config` command to create a config file in `~/.config/llmscript/config.yaml` which you can edit. You can also use command-line args (see below). 91 | 92 | ## How it works 93 | 94 | Want to see it all in action? Run `llmscript --verbose examples/hello-world` 95 | 96 | Given a script description written in natural language, llmscript works by: 97 | 98 | 1. Generating a feature script that implements the functionality 99 | 2. Generating a test script that verifies the feature script works correctly 100 | 3. Running the test script to verify the feature script works correctly, fixing the feature script if necessary, possibly going back to step 1 if the test script fails too many times 101 | 4. Caching the scripts for future use 102 | 5. Running the feature script with any additional arguments you provide 103 | 104 | For example, given a simple hello world script: 105 | 106 | ``` 107 | #!/usr/bin/env llmscript 108 | 109 | Print hello world 110 | ``` 111 | 112 | llmscript might generate the following feature script: 113 | 114 | ```bash 115 | #!/bin/bash 116 | echo "Hello, world!" 117 | ``` 118 | 119 | ...and the following test script to test it: 120 | 121 | ```bash 122 | #!/bin/bash 123 | [ "$(./script.sh)" = "Hello, world!" ] || exit 1 124 | ``` 125 | 126 | ## Configuration 127 | 128 | llmscript can be configured using a YAML file located at `~/.config/llmscript/config.yaml`. You can auto-generate a configuration file using the `llmscript --write-config` command. 129 | 130 | Here's an example configuration: 131 | 132 | ```yaml 133 | # LLM configuration 134 | llm: 135 | # The LLM provider to use (required) 136 | provider: "ollama" # or "claude", "openai", etc. 137 | 138 | # Provider-specific settings 139 | ollama: 140 | model: "llama3.2" # The model to use 141 | host: "http://localhost:11434" # Optional: Ollama host URL 142 | 143 | claude: 144 | api_key: "${CLAUDE_API_KEY}" # Environment variable reference 145 | model: "claude-3-opus-20240229" 146 | 147 | openai: 148 | api_key: "${OPENAI_API_KEY}" 149 | model: "gpt-4-turbo-preview" 150 | 151 | # Maximum number of attempts to fix the script allowed before restarting from step 2 152 | max_fixes: 10 153 | 154 | # Maximum number of attempts to generate a working script before giving up completely 155 | max_attempts: 3 156 | 157 | # Timeout for script execution during testing (in seconds) 158 | timeout: 30 159 | 160 | # Additional prompt to provide to the LLM 161 | additional_prompt: | 162 | Use ANSI color codes to make the output more readable. 163 | ``` 164 | 165 | ### Environment Variables 166 | 167 | You can use environment variables in the configuration file using the `${VAR_NAME}` syntax. This is particularly useful for API keys and sensitive information. 168 | 169 | ### Configuration Precedence 170 | 171 | 1. Command line flags (highest priority) 172 | 2. Environment variables 173 | 3. Configuration file 174 | 4. Default values (lowest priority) 175 | 176 | ### Command Line Flags 177 | 178 | You can override configuration settings using command line flags: 179 | 180 | ```shell 181 | llmscript --llm.provider=claude --timeout=10 script.txt 182 | ``` 183 | 184 | ## Caveats 185 | 186 | > [!WARNING] 187 | > This is an experimental project. It generates and executes shell scripts, which could be dangerous if the LLM generates malicious code. Use at your own risk and always review generated scripts before running them. 188 | -------------------------------------------------------------------------------- /internal/llm/claude.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "fmt" 7 | "io" 8 | "net/http" 9 | "strings" 10 | 11 | "github.com/statico/llmscript/internal/log" 12 | ) 13 | 14 | // ClaudeProvider implements the Provider interface using Anthropic's Claude 15 | type ClaudeProvider struct { 16 | BaseProvider 17 | config ClaudeConfig 18 | } 19 | 20 | // NewClaudeProvider creates a new Claude provider 21 | func NewClaudeProvider(config ClaudeConfig) (*ClaudeProvider, error) { 22 | return &ClaudeProvider{ 23 | BaseProvider: BaseProvider{Config: config}, 24 | config: config, 25 | }, nil 26 | } 27 | 28 | // GenerateScripts creates a main script and test script from a natural language description 29 | func (p *ClaudeProvider) GenerateScripts(ctx context.Context, description string) (ScriptPair, error) { 30 | // First generate the main script 31 | log.Info("Generating main script with Claude...") 32 | mainPrompt := p.formatPrompt(FeatureScriptPrompt, description) 33 | mainScript, err := p.generate(ctx, mainPrompt) 34 | if err != nil { 35 | return ScriptPair{}, fmt.Errorf("failed to generate main script: %w", err) 36 | } 37 | mainScript = ExtractScriptContent(mainScript) 38 | log.Debug("Main script generated:\n%s", mainScript) 39 | 40 | // Then generate the test script 41 | log.Info("Generating test script with Claude...") 42 | testPrompt := p.formatPrompt(TestScriptPrompt, mainScript, description) 43 | testScript, err := p.generate(ctx, testPrompt) 44 | if err != nil { 45 | return ScriptPair{}, fmt.Errorf("failed to generate test script: %w", err) 46 | } 47 | testScript = ExtractScriptContent(testScript) 48 | log.Debug("Test script generated:\n%s", testScript) 49 | 50 | return ScriptPair{ 51 | MainScript: strings.TrimSpace(mainScript), 52 | TestScript: strings.TrimSpace(testScript), 53 | }, nil 54 | } 55 | 56 | // FixScripts attempts to fix both scripts based on test failures 57 | func (p *ClaudeProvider) FixScripts(ctx context.Context, scripts ScriptPair, error string) (ScriptPair, error) { 58 | // Only fix the main script 59 | mainPrompt := p.formatPrompt(FixScriptPrompt, scripts.MainScript, error) 60 | fixedMainScript, err := p.generate(ctx, mainPrompt) 61 | if err != nil { 62 | return ScriptPair{}, fmt.Errorf("failed to fix main script: %w", err) 63 | } 64 | fixedMainScript = ExtractScriptContent(fixedMainScript) 65 | 66 | return ScriptPair{ 67 | MainScript: strings.TrimSpace(fixedMainScript), 68 | TestScript: scripts.TestScript, 69 | }, nil 70 | } 71 | 72 | // GenerateScript creates a shell script from a natural language description 73 | func (p *ClaudeProvider) GenerateScript(ctx context.Context, description string) (string, error) { 74 | prompt := p.formatPrompt(FeatureScriptPrompt, description) 75 | response, err := p.generate(ctx, prompt) 76 | if err != nil { 77 | return "", fmt.Errorf("failed to generate script: %w", err) 78 | } 79 | return response, nil 80 | } 81 | 82 | // GenerateTests creates test cases for a script based on its description 83 | func (p *ClaudeProvider) GenerateTests(ctx context.Context, script string, description string) ([]Test, error) { 84 | prompt := p.formatPrompt(TestScriptPrompt, script, description) 85 | response, err := p.generate(ctx, prompt) 86 | if err != nil { 87 | return nil, fmt.Errorf("failed to generate tests: %w", err) 88 | } 89 | 90 | log.Debug("Raw LLM response:\n%s", response) 91 | 92 | // Try to extract JSON from the response 93 | jsonStart := strings.Index(response, "{") 94 | jsonEnd := strings.LastIndex(response, "}") 95 | if jsonStart == -1 || jsonEnd == -1 || jsonEnd <= jsonStart { 96 | return nil, fmt.Errorf("failed to find valid JSON in response: %s", response) 97 | } 98 | jsonStr := response[jsonStart : jsonEnd+1] 99 | 100 | var result struct { 101 | Tests []Test `json:"tests"` 102 | } 103 | if err := json.Unmarshal([]byte(jsonStr), &result); err != nil { 104 | return nil, fmt.Errorf("failed to parse test cases: %w\nRaw response:\n%s", err, response) 105 | } 106 | return result.Tests, nil 107 | } 108 | 109 | // FixScript attempts to fix a script based on test failures 110 | func (p *ClaudeProvider) FixScript(ctx context.Context, script string, failures []TestFailure) (string, error) { 111 | failuresStr := formatFailures(failures) 112 | prompt := p.formatPrompt(FixScriptPrompt, script, failuresStr) 113 | response, err := p.generate(ctx, prompt) 114 | if err != nil { 115 | return "", fmt.Errorf("failed to fix script: %w", err) 116 | } 117 | return response, nil 118 | } 119 | 120 | // generate sends a prompt to Claude and returns the response 121 | func (p *ClaudeProvider) generate(ctx context.Context, prompt string) (string, error) { 122 | url := "https://api.anthropic.com/v1/messages" 123 | 124 | reqBody := map[string]interface{}{ 125 | "model": p.config.Model, 126 | "messages": []map[string]string{{"role": "user", "content": prompt}}, 127 | "max_tokens": 4096, 128 | } 129 | 130 | jsonData, err := json.Marshal(reqBody) 131 | if err != nil { 132 | return "", fmt.Errorf("failed to marshal request: %w", err) 133 | } 134 | 135 | req, err := http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(jsonData))) 136 | if err != nil { 137 | return "", fmt.Errorf("failed to create request: %w", err) 138 | } 139 | req.Header.Set("Content-Type", "application/json") 140 | req.Header.Set("x-api-key", p.config.APIKey) 141 | req.Header.Set("anthropic-version", "2023-06-01") 142 | 143 | resp, err := http.DefaultClient.Do(req) 144 | if err != nil { 145 | return "", fmt.Errorf("failed to send request: %w", err) 146 | } 147 | defer func() { 148 | if err := resp.Body.Close(); err != nil { 149 | log.Error("failed to close response body: %v", err) 150 | } 151 | }() 152 | 153 | if resp.StatusCode != http.StatusOK { 154 | body, _ := io.ReadAll(resp.Body) 155 | return "", fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(body)) 156 | } 157 | 158 | var result struct { 159 | Content []struct { 160 | Text string `json:"text"` 161 | } `json:"content"` 162 | } 163 | if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { 164 | return "", fmt.Errorf("failed to decode response: %w", err) 165 | } 166 | 167 | if len(result.Content) == 0 { 168 | return "", fmt.Errorf("no content in response") 169 | } 170 | 171 | return result.Content[0].Text, nil 172 | } 173 | 174 | // Name returns a human-readable name for the provider 175 | func (p *ClaudeProvider) Name() string { 176 | return "Claude" 177 | } 178 | -------------------------------------------------------------------------------- /internal/llm/openai.go: -------------------------------------------------------------------------------- 1 | package llm 2 | 3 | import ( 4 | "context" 5 | "encoding/json" 6 | "fmt" 7 | "io" 8 | "net/http" 9 | "strings" 10 | 11 | "github.com/statico/llmscript/internal/log" 12 | ) 13 | 14 | // OpenAIProvider implements the Provider interface using OpenAI's API 15 | type OpenAIProvider struct { 16 | BaseProvider 17 | config OpenAIConfig 18 | } 19 | 20 | // NewOpenAIProvider creates a new OpenAI provider 21 | func NewOpenAIProvider(config OpenAIConfig) (*OpenAIProvider, error) { 22 | return &OpenAIProvider{ 23 | BaseProvider: BaseProvider{Config: config}, 24 | config: config, 25 | }, nil 26 | } 27 | 28 | // GenerateScripts creates a main script and test script from a natural language description 29 | func (p *OpenAIProvider) GenerateScripts(ctx context.Context, description string) (ScriptPair, error) { 30 | // First generate the main script 31 | log.Info("Generating main script with OpenAI...") 32 | mainPrompt := p.formatPrompt(FeatureScriptPrompt, description) 33 | mainScript, err := p.generate(ctx, mainPrompt) 34 | if err != nil { 35 | return ScriptPair{}, fmt.Errorf("failed to generate main script: %w", err) 36 | } 37 | mainScript = ExtractScriptContent(mainScript) 38 | log.Debug("Main script generated:\n%s", mainScript) 39 | 40 | // Then generate the test script 41 | log.Info("Generating test script with OpenAI...") 42 | testPrompt := p.formatPrompt(TestScriptPrompt, mainScript, description) 43 | testScript, err := p.generate(ctx, testPrompt) 44 | if err != nil { 45 | return ScriptPair{}, fmt.Errorf("failed to generate test script: %w", err) 46 | } 47 | testScript = ExtractScriptContent(testScript) 48 | log.Debug("Test script generated:\n%s", testScript) 49 | 50 | return ScriptPair{ 51 | MainScript: strings.TrimSpace(mainScript), 52 | TestScript: strings.TrimSpace(testScript), 53 | }, nil 54 | } 55 | 56 | // FixScripts attempts to fix both scripts based on test failures 57 | func (p *OpenAIProvider) FixScripts(ctx context.Context, scripts ScriptPair, error string) (ScriptPair, error) { 58 | // Only fix the main script 59 | mainPrompt := p.formatPrompt(FixScriptPrompt, scripts.MainScript, error) 60 | fixedMainScript, err := p.generate(ctx, mainPrompt) 61 | if err != nil { 62 | return ScriptPair{}, fmt.Errorf("failed to fix main script: %w", err) 63 | } 64 | fixedMainScript = ExtractScriptContent(fixedMainScript) 65 | 66 | return ScriptPair{ 67 | MainScript: strings.TrimSpace(fixedMainScript), 68 | TestScript: scripts.TestScript, 69 | }, nil 70 | } 71 | 72 | // GenerateScript creates a shell script from a natural language description 73 | func (p *OpenAIProvider) GenerateScript(ctx context.Context, description string) (string, error) { 74 | prompt := p.formatPrompt(FeatureScriptPrompt, description) 75 | response, err := p.generate(ctx, prompt) 76 | if err != nil { 77 | return "", fmt.Errorf("failed to generate script: %w", err) 78 | } 79 | return response, nil 80 | } 81 | 82 | // GenerateTests creates test cases for a script based on its description 83 | func (p *OpenAIProvider) GenerateTests(ctx context.Context, script string, description string) ([]Test, error) { 84 | prompt := p.formatPrompt(TestScriptPrompt, script, description) 85 | response, err := p.generate(ctx, prompt) 86 | if err != nil { 87 | return nil, fmt.Errorf("failed to generate tests: %w", err) 88 | } 89 | 90 | log.Debug("Raw LLM response:\n%s", response) 91 | 92 | // Try to extract JSON from the response 93 | jsonStart := strings.Index(response, "{") 94 | jsonEnd := strings.LastIndex(response, "}") 95 | if jsonStart == -1 || jsonEnd == -1 || jsonEnd <= jsonStart { 96 | return nil, fmt.Errorf("failed to find valid JSON in response: %s", response) 97 | } 98 | jsonStr := response[jsonStart : jsonEnd+1] 99 | 100 | var result struct { 101 | Tests []Test `json:"tests"` 102 | } 103 | if err := json.Unmarshal([]byte(jsonStr), &result); err != nil { 104 | return nil, fmt.Errorf("failed to parse test cases: %w\nRaw response:\n%s", err, response) 105 | } 106 | return result.Tests, nil 107 | } 108 | 109 | // FixScript attempts to fix a script based on test failures 110 | func (p *OpenAIProvider) FixScript(ctx context.Context, script string, failures []TestFailure) (string, error) { 111 | failuresStr := formatFailures(failures) 112 | prompt := p.formatPrompt(FixScriptPrompt, script, failuresStr) 113 | response, err := p.generate(ctx, prompt) 114 | if err != nil { 115 | return "", fmt.Errorf("failed to fix script: %w", err) 116 | } 117 | return response, nil 118 | } 119 | 120 | // generate sends a prompt to OpenAI and returns the response 121 | func (p *OpenAIProvider) generate(ctx context.Context, prompt string) (string, error) { 122 | url := "https://api.openai.com/v1/chat/completions" 123 | 124 | reqBody := map[string]interface{}{ 125 | "model": p.config.Model, 126 | "messages": []map[string]string{ 127 | {"role": "user", "content": prompt}, 128 | }, 129 | } 130 | 131 | jsonData, err := json.Marshal(reqBody) 132 | if err != nil { 133 | return "", fmt.Errorf("failed to marshal request: %w", err) 134 | } 135 | 136 | req, err := http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(jsonData))) 137 | if err != nil { 138 | return "", fmt.Errorf("failed to create request: %w", err) 139 | } 140 | req.Header.Set("Content-Type", "application/json") 141 | req.Header.Set("Authorization", "Bearer "+p.config.APIKey) 142 | 143 | resp, err := http.DefaultClient.Do(req) 144 | if err != nil { 145 | return "", fmt.Errorf("failed to send request: %w", err) 146 | } 147 | defer func() { 148 | if err := resp.Body.Close(); err != nil { 149 | log.Error("failed to close response body: %v", err) 150 | } 151 | }() 152 | 153 | if resp.StatusCode != http.StatusOK { 154 | body, _ := io.ReadAll(resp.Body) 155 | return "", fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(body)) 156 | } 157 | 158 | var result struct { 159 | Choices []struct { 160 | Message struct { 161 | Content string `json:"content"` 162 | } `json:"message"` 163 | } `json:"choices"` 164 | } 165 | if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { 166 | return "", fmt.Errorf("failed to decode response: %w", err) 167 | } 168 | 169 | if len(result.Choices) == 0 { 170 | return "", fmt.Errorf("no choices in response") 171 | } 172 | 173 | return result.Choices[0].Message.Content, nil 174 | } 175 | 176 | // Name returns a human-readable name for the provider 177 | func (p *OpenAIProvider) Name() string { 178 | return "OpenAI" 179 | } 180 | -------------------------------------------------------------------------------- /internal/config/config.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "fmt" 5 | "log" 6 | "os" 7 | "path/filepath" 8 | "time" 9 | 10 | customlog "github.com/statico/llmscript/internal/log" 11 | "gopkg.in/yaml.v3" 12 | ) 13 | 14 | type Config struct { 15 | LLM struct { 16 | Provider string `yaml:"provider"` 17 | Ollama struct { 18 | Model string `yaml:"model"` 19 | Host string `yaml:"host"` 20 | } `yaml:"ollama"` 21 | Claude struct { 22 | APIKey string `yaml:"api_key"` 23 | Model string `yaml:"model"` 24 | } `yaml:"claude"` 25 | OpenAI struct { 26 | APIKey string `yaml:"api_key"` 27 | Model string `yaml:"model"` 28 | } `yaml:"openai"` 29 | } `yaml:"llm"` 30 | MaxFixes int `yaml:"max_fixes"` 31 | MaxAttempts int `yaml:"max_attempts"` 32 | Timeout time.Duration `yaml:"timeout"` 33 | ExtraPrompt string `yaml:"additional_prompt"` 34 | } 35 | 36 | func DefaultConfig() *Config { 37 | return &Config{ 38 | LLM: struct { 39 | Provider string `yaml:"provider"` 40 | Ollama struct { 41 | Model string `yaml:"model"` 42 | Host string `yaml:"host"` 43 | } `yaml:"ollama"` 44 | Claude struct { 45 | APIKey string `yaml:"api_key"` 46 | Model string `yaml:"model"` 47 | } `yaml:"claude"` 48 | OpenAI struct { 49 | APIKey string `yaml:"api_key"` 50 | Model string `yaml:"model"` 51 | } `yaml:"openai"` 52 | }{ 53 | Provider: "ollama", 54 | Ollama: struct { 55 | Model string `yaml:"model"` 56 | Host string `yaml:"host"` 57 | }{ 58 | Model: "llama3.2", 59 | Host: "http://localhost:11434", 60 | }, 61 | Claude: struct { 62 | APIKey string `yaml:"api_key"` 63 | Model string `yaml:"model"` 64 | }{ 65 | APIKey: "${CLAUDE_API_KEY}", 66 | Model: "claude-3-opus-20240229", 67 | }, 68 | OpenAI: struct { 69 | APIKey string `yaml:"api_key"` 70 | Model string `yaml:"model"` 71 | }{ 72 | APIKey: "${OPENAI_API_KEY}", 73 | Model: "gpt-4-turbo-preview", 74 | }, 75 | }, 76 | MaxFixes: 10, 77 | MaxAttempts: 3, 78 | Timeout: 30 * time.Second, 79 | ExtraPrompt: "Use ANSI color codes to make the output more readable.", 80 | } 81 | } 82 | 83 | func interpolateEnvVars(data []byte) []byte { 84 | content := string(data) 85 | content = os.ExpandEnv(content) 86 | return []byte(content) 87 | } 88 | 89 | func LoadConfig() (*Config, error) { 90 | config := DefaultConfig() 91 | 92 | // Try XDG_CONFIG_HOME first 93 | configDir := os.Getenv("XDG_CONFIG_HOME") 94 | if configDir == "" { 95 | // If XDG_CONFIG_HOME is not set, try ~/.config 96 | homeDir, err := os.UserHomeDir() 97 | if err != nil { 98 | return nil, fmt.Errorf("failed to get home directory: %w", err) 99 | } 100 | configDir = filepath.Join(homeDir, ".config") 101 | } 102 | 103 | configPath := filepath.Join(configDir, "llmscript", "config.yaml") 104 | customlog.Debug("Looking for config file at: %s", configPath) 105 | data, err := os.ReadFile(configPath) 106 | if err != nil { 107 | if os.IsNotExist(err) { 108 | // If not found in ~/.config, try UserConfigDir() as fallback 109 | configDir, err = os.UserConfigDir() 110 | if err != nil { 111 | return nil, fmt.Errorf("failed to get config dir: %w", err) 112 | } 113 | configPath = filepath.Join(configDir, "llmscript", "config.yaml") 114 | customlog.Debug("Looking for config file at fallback location: %s", configPath) 115 | data, err = os.ReadFile(configPath) 116 | if err != nil { 117 | if os.IsNotExist(err) { 118 | customlog.Debug("Config file not found, using defaults") 119 | return config, nil 120 | } 121 | return nil, fmt.Errorf("failed to read config file: %w", err) 122 | } 123 | } else { 124 | return nil, fmt.Errorf("failed to read config file: %w", err) 125 | } 126 | } 127 | 128 | customlog.Debug("Found config file: %s", configPath) 129 | 130 | // Interpolate environment variables before unmarshaling 131 | data = interpolateEnvVars(data) 132 | 133 | // Create a temporary config to unmarshal into 134 | var loadedConfig Config 135 | if err := yaml.Unmarshal(data, &loadedConfig); err != nil { 136 | return nil, fmt.Errorf("failed to parse config file: %w", err) 137 | } 138 | 139 | customlog.Debug("Loaded config from file: provider=%s", loadedConfig.LLM.Provider) 140 | 141 | // Merge the loaded config with the default config 142 | config.LLM.Provider = loadedConfig.LLM.Provider 143 | config.LLM.Ollama.Model = loadedConfig.LLM.Ollama.Model 144 | config.LLM.Ollama.Host = loadedConfig.LLM.Ollama.Host 145 | config.LLM.Claude.APIKey = loadedConfig.LLM.Claude.APIKey 146 | config.LLM.Claude.Model = loadedConfig.LLM.Claude.Model 147 | config.LLM.OpenAI.APIKey = loadedConfig.LLM.OpenAI.APIKey 148 | config.LLM.OpenAI.Model = loadedConfig.LLM.OpenAI.Model 149 | 150 | customlog.Debug("Merged config: provider=%s", config.LLM.Provider) 151 | 152 | // Only update numeric values if they are explicitly set in the YAML 153 | if loadedConfig.MaxFixes != 0 { 154 | config.MaxFixes = loadedConfig.MaxFixes 155 | } 156 | if loadedConfig.MaxAttempts != 0 { 157 | config.MaxAttempts = loadedConfig.MaxAttempts 158 | } 159 | if loadedConfig.Timeout != 0 { 160 | config.Timeout = loadedConfig.Timeout 161 | } 162 | 163 | config.ExtraPrompt = loadedConfig.ExtraPrompt 164 | 165 | return config, nil 166 | } 167 | 168 | func WriteConfig(config *Config) error { 169 | configDir := os.Getenv("XDG_CONFIG_HOME") 170 | if configDir == "" { 171 | homeDir, err := os.UserHomeDir() 172 | if err != nil { 173 | return fmt.Errorf("failed to get home directory: %w", err) 174 | } 175 | configDir = filepath.Join(homeDir, ".config") 176 | } 177 | customlog.Debug("Config directory: %s", configDir) 178 | 179 | llmscriptDir := filepath.Join(configDir, "llmscript") 180 | customlog.Debug("Creating directory: %s", llmscriptDir) 181 | if err := os.MkdirAll(llmscriptDir, 0755); err != nil { 182 | return fmt.Errorf("failed to create config directory: %w", err) 183 | } 184 | 185 | configPath := filepath.Join(llmscriptDir, "config.yaml") 186 | customlog.Debug("Writing config to: %s", configPath) 187 | f, err := os.Create(configPath) 188 | if err != nil { 189 | return fmt.Errorf("failed to create config file: %w", err) 190 | } 191 | defer func() { 192 | if err := f.Close(); err != nil { 193 | log.Printf("failed to close config file: %v", err) 194 | } 195 | }() 196 | 197 | encoder := yaml.NewEncoder(f) 198 | encoder.SetIndent(2) 199 | if err := encoder.Encode(config); err != nil { 200 | return fmt.Errorf("failed to encode config: %w", err) 201 | } 202 | if err := encoder.Close(); err != nil { 203 | return fmt.Errorf("failed to close encoder: %w", err) 204 | } 205 | 206 | return nil 207 | } 208 | -------------------------------------------------------------------------------- /cmd/llmscript/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "context" 5 | "flag" 6 | "fmt" 7 | "os" 8 | "os/exec" 9 | "path/filepath" 10 | "time" 11 | 12 | "github.com/statico/llmscript/internal/config" 13 | "github.com/statico/llmscript/internal/llm" 14 | "github.com/statico/llmscript/internal/log" 15 | "github.com/statico/llmscript/internal/script" 16 | ) 17 | 18 | var ( 19 | writeConfig = flag.Bool("write-config", false, "Write default config to ~/.config/llmscript/config.yaml") 20 | verbose = flag.Bool("verbose", false, "Enable verbose output (includes debug messages)") 21 | timeout = flag.Duration("timeout", 30*time.Second, "Timeout for script execution") 22 | maxFixes = flag.Int("max-fixes", 3, "Maximum number of attempts to fix the script") 23 | maxAttempts = flag.Int("max-attempts", 3, "Maximum number of attempts to generate a working script") 24 | llmProvider = flag.String("llm.provider", "", "LLM provider to use (overrides config)") 25 | llmModel = flag.String("llm.model", "", "LLM model to use (overrides config)") 26 | extraPrompt = flag.String("prompt", "", "Additional prompt to provide to the LLM") 27 | noCache = flag.Bool("no-cache", false, "Skip using the cache for script generation") 28 | printOnly = flag.Bool("print", false, "Print the generated script without executing it") 29 | ) 30 | 31 | func main() { 32 | flag.Usage = func() { 33 | fmt.Fprintf(os.Stderr, "Usage: %s [flags] \n\n", os.Args[0]) 34 | fmt.Fprintf(os.Stderr, "Flags:\n") 35 | flag.PrintDefaults() 36 | fmt.Fprintf(os.Stderr, "\nExamples:\n") 37 | fmt.Fprintf(os.Stderr, " %s script.txt\n", os.Args[0]) 38 | fmt.Fprintf(os.Stderr, " %s --llm.provider=claude --timeout=10 script.txt\n", os.Args[0]) 39 | fmt.Fprintf(os.Stderr, " %s --write-config\n", os.Args[0]) 40 | } 41 | flag.Parse() 42 | 43 | // Set up logging first 44 | if *verbose { 45 | log.SetLevel(log.DebugLevel) 46 | } else { 47 | log.SetLevel(log.InfoLevel) 48 | } 49 | 50 | if *writeConfig { 51 | cfg := config.DefaultConfig() 52 | if err := config.WriteConfig(cfg); err != nil { 53 | log.Fatal("Failed to write default config:", err) 54 | } 55 | fmt.Println("Default config written to ~/.config/llmscript/config.yaml") 56 | return 57 | } 58 | 59 | cfg, err := config.LoadConfig() 60 | if err != nil { 61 | log.Fatal("Failed to load config:", err) 62 | } 63 | 64 | log.Debug("Loaded config: provider=%s, has_claude_api_key=%v, claude_model=%s", 65 | cfg.LLM.Provider, 66 | cfg.LLM.Claude.APIKey != "", 67 | cfg.LLM.Claude.Model) 68 | 69 | // Override config with command line flags 70 | if *llmProvider != "" { 71 | cfg.LLM.Provider = *llmProvider 72 | log.Debug("Provider overridden by command line flag: %s", *llmProvider) 73 | } 74 | if *llmModel != "" { 75 | switch cfg.LLM.Provider { 76 | case "ollama": 77 | cfg.LLM.Ollama.Model = *llmModel 78 | case "claude": 79 | cfg.LLM.Claude.Model = *llmModel 80 | case "openai": 81 | cfg.LLM.OpenAI.Model = *llmModel 82 | } 83 | } 84 | if *timeout != 0 { 85 | cfg.Timeout = *timeout 86 | } 87 | if *maxFixes != 0 { 88 | cfg.MaxFixes = *maxFixes 89 | } 90 | if *maxAttempts != 0 { 91 | cfg.MaxAttempts = *maxAttempts 92 | } 93 | if *extraPrompt != "" { 94 | cfg.ExtraPrompt = *extraPrompt 95 | } 96 | 97 | if len(flag.Args()) == 0 { 98 | flag.Usage() 99 | os.Exit(1) 100 | } 101 | 102 | scriptFile := flag.Args()[0] 103 | if err := runScript(cfg, scriptFile); err != nil { 104 | log.Fatal("Failed to run script:", err) 105 | } 106 | } 107 | 108 | func runScript(cfg *config.Config, scriptFile string) error { 109 | log.Info("Reading script file: %s", scriptFile) 110 | content, err := os.ReadFile(scriptFile) 111 | if err != nil { 112 | return fmt.Errorf("failed to read script file: %w", err) 113 | } 114 | 115 | log.Info("Creating LLM provider: %s", cfg.LLM.Provider) 116 | provider, err := llm.NewProvider(cfg.LLM.Provider, map[string]interface{}{ 117 | "ollama": map[string]interface{}{ 118 | "model": cfg.LLM.Ollama.Model, 119 | "host": cfg.LLM.Ollama.Host, 120 | }, 121 | "claude": map[string]interface{}{ 122 | "api_key": cfg.LLM.Claude.APIKey, 123 | "model": cfg.LLM.Claude.Model, 124 | }, 125 | "openai": map[string]interface{}{ 126 | "api_key": cfg.LLM.OpenAI.APIKey, 127 | "model": cfg.LLM.OpenAI.Model, 128 | }, 129 | }) 130 | if err != nil { 131 | return fmt.Errorf("failed to create LLM provider: %w", err) 132 | } 133 | 134 | log.Info("Creating work directory") 135 | workDir, err := os.MkdirTemp("", "llmscript-*") 136 | if err != nil { 137 | return fmt.Errorf("failed to create working directory: %w", err) 138 | } 139 | defer func() { 140 | if err := os.RemoveAll(workDir); err != nil { 141 | log.Error("failed to remove working directory: %v", err) 142 | } 143 | }() 144 | if *verbose { 145 | log.Info("Work directory: %s", workDir) 146 | } 147 | 148 | log.Info("Creating pipeline") 149 | pipeline, err := script.NewPipeline(provider, cfg.MaxFixes, cfg.MaxAttempts, cfg.Timeout, workDir, *noCache) 150 | if err != nil { 151 | return fmt.Errorf("failed to create pipeline: %w", err) 152 | } 153 | 154 | ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout) 155 | defer cancel() 156 | 157 | log.Info("Generating and testing script") 158 | script, err := pipeline.GenerateAndTest(ctx, string(content)) 159 | if err != nil { 160 | return fmt.Errorf("failed to generate working script: %w", err) 161 | } 162 | 163 | if *verbose { 164 | log.Info("Generated script:\n%s", script) 165 | } 166 | 167 | // Write the script to a file 168 | scriptPath := filepath.Join(workDir, "script.sh") 169 | if err := os.WriteFile(scriptPath, []byte(script), 0755); err != nil { 170 | return fmt.Errorf("failed to write script: %w", err) 171 | } 172 | 173 | // Clear the spinner line before printing success message 174 | log.GetSpinner().Clear() 175 | 176 | // If --print flag is set, just print the script and exit 177 | if *printOnly { 178 | fmt.Println(script) 179 | return nil 180 | } 181 | 182 | // Stop the spinner before executing the script 183 | log.GetSpinner().Stop() 184 | 185 | // Get any additional arguments after the script file 186 | scriptArgs := flag.Args()[1:] 187 | 188 | // Execute the script with any additional arguments 189 | cmd := exec.Command(scriptPath, scriptArgs...) 190 | cmd.Stdout = os.Stdout 191 | cmd.Stderr = os.Stderr 192 | cmd.Stdin = os.Stdin 193 | 194 | // Run the command and exit with its status code 195 | if err := cmd.Run(); err != nil { 196 | if exitErr, ok := err.(*exec.ExitError); ok { 197 | os.Exit(exitErr.ExitCode()) 198 | } 199 | // If it's not an exit error, something else went wrong 200 | return fmt.Errorf("script execution failed: %w", err) 201 | } 202 | 203 | return nil 204 | } 205 | -------------------------------------------------------------------------------- /internal/config/config_test.go: -------------------------------------------------------------------------------- 1 | package config 2 | 3 | import ( 4 | "os" 5 | "path/filepath" 6 | "testing" 7 | "time" 8 | ) 9 | 10 | func TestConfig(t *testing.T) { 11 | t.Run("default config", func(t *testing.T) { 12 | cfg := DefaultConfig() 13 | if cfg.MaxFixes != 10 { 14 | t.Errorf("expected MaxFixes=10, got %d", cfg.MaxFixes) 15 | } 16 | if cfg.MaxAttempts != 3 { 17 | t.Errorf("expected MaxAttempts=3, got %d", cfg.MaxAttempts) 18 | } 19 | if cfg.Timeout != 30*time.Second { 20 | t.Errorf("expected Timeout=30s, got %v", cfg.Timeout) 21 | } 22 | }) 23 | 24 | t.Run("load and write config", func(t *testing.T) { 25 | tmpDir := t.TempDir() 26 | if err := os.Setenv("XDG_CONFIG_HOME", tmpDir); err != nil { 27 | t.Fatalf("failed to set XDG_CONFIG_HOME: %v", err) 28 | } 29 | defer func() { 30 | if err := os.Unsetenv("XDG_CONFIG_HOME"); err != nil { 31 | t.Errorf("failed to unset XDG_CONFIG_HOME: %v", err) 32 | } 33 | }() 34 | 35 | cfg := DefaultConfig() 36 | cfg.LLM.Provider = "ollama" 37 | cfg.LLM.Ollama.Model = "llama2" 38 | cfg.LLM.Ollama.Host = "http://localhost:11434" 39 | 40 | if err := WriteConfig(cfg); err != nil { 41 | t.Fatalf("failed to write config: %v", err) 42 | } 43 | 44 | loaded, err := LoadConfig() 45 | if err != nil { 46 | t.Fatalf("failed to load config: %v", err) 47 | } 48 | 49 | if loaded.LLM.Provider != cfg.LLM.Provider { 50 | t.Errorf("expected provider=%s, got %s", cfg.LLM.Provider, loaded.LLM.Provider) 51 | } 52 | if loaded.LLM.Ollama.Model != cfg.LLM.Ollama.Model { 53 | t.Errorf("expected model=%s, got %s", cfg.LLM.Ollama.Model, loaded.LLM.Ollama.Model) 54 | } 55 | if loaded.LLM.Ollama.Host != cfg.LLM.Ollama.Host { 56 | t.Errorf("expected host=%s, got %s", cfg.LLM.Ollama.Host, loaded.LLM.Ollama.Host) 57 | } 58 | }) 59 | 60 | t.Run("load non-existent config", func(t *testing.T) { 61 | tmpDir := t.TempDir() 62 | if err := os.Setenv("XDG_CONFIG_HOME", tmpDir); err != nil { 63 | t.Fatalf("failed to set XDG_CONFIG_HOME: %v", err) 64 | } 65 | defer func() { 66 | if err := os.Unsetenv("XDG_CONFIG_HOME"); err != nil { 67 | t.Errorf("failed to unset XDG_CONFIG_HOME: %v", err) 68 | } 69 | }() 70 | 71 | cfg, err := LoadConfig() 72 | if err != nil { 73 | t.Fatalf("failed to load config: %v", err) 74 | } 75 | 76 | if cfg.MaxFixes != 10 { 77 | t.Errorf("expected MaxFixes=10, got %d", cfg.MaxFixes) 78 | } 79 | }) 80 | 81 | t.Run("load yaml file with env vars", func(t *testing.T) { 82 | tmpDir := t.TempDir() 83 | if err := os.Setenv("XDG_CONFIG_HOME", tmpDir); err != nil { 84 | t.Fatalf("failed to set XDG_CONFIG_HOME: %v", err) 85 | } 86 | defer func() { 87 | if err := os.Unsetenv("XDG_CONFIG_HOME"); err != nil { 88 | t.Errorf("failed to unset XDG_CONFIG_HOME: %v", err) 89 | } 90 | }() 91 | 92 | // Set up test environment variables 93 | if err := os.Setenv("TEST_API_KEY", "test-key-123"); err != nil { 94 | t.Fatalf("failed to set TEST_API_KEY: %v", err) 95 | } 96 | defer func() { 97 | if err := os.Unsetenv("TEST_API_KEY"); err != nil { 98 | t.Errorf("failed to unset TEST_API_KEY: %v", err) 99 | } 100 | }() 101 | 102 | // Create test config directory 103 | configDir := filepath.Join(tmpDir, "llmscript") 104 | if err := os.MkdirAll(configDir, 0755); err != nil { 105 | t.Fatalf("failed to create config directory: %v", err) 106 | } 107 | 108 | // Create a test YAML file 109 | yamlContent := []byte(`llm: 110 | provider: "claude" 111 | claude: 112 | api_key: "${TEST_API_KEY}" 113 | model: "claude-3-opus-20240229" 114 | max_fixes: 5 115 | max_attempts: 2 116 | timeout: 15s 117 | additional_prompt: "Test prompt" 118 | `) 119 | configPath := filepath.Join(configDir, "config.yaml") 120 | if err := os.WriteFile(configPath, yamlContent, 0644); err != nil { 121 | t.Fatalf("failed to write test config: %v", err) 122 | } 123 | 124 | cfg, err := LoadConfig() 125 | if err != nil { 126 | t.Fatalf("failed to load config: %v", err) 127 | } 128 | 129 | if cfg.LLM.Provider != "claude" { 130 | t.Errorf("expected provider=claude, got %s", cfg.LLM.Provider) 131 | } 132 | if cfg.LLM.Claude.APIKey != "test-key-123" { 133 | t.Errorf("expected api_key=test-key-123, got %s", cfg.LLM.Claude.APIKey) 134 | } 135 | if cfg.LLM.Claude.Model != "claude-3-opus-20240229" { 136 | t.Errorf("expected model=claude-3-opus-20240229, got %s", cfg.LLM.Claude.Model) 137 | } 138 | if cfg.MaxFixes != 5 { 139 | t.Errorf("expected MaxFixes=5, got %d", cfg.MaxFixes) 140 | } 141 | if cfg.MaxAttempts != 2 { 142 | t.Errorf("expected MaxAttempts=2, got %d", cfg.MaxAttempts) 143 | } 144 | if cfg.Timeout != 15*time.Second { 145 | t.Errorf("expected Timeout=15s, got %v", cfg.Timeout) 146 | } 147 | if cfg.ExtraPrompt != "Test prompt" { 148 | t.Errorf("expected ExtraPrompt=Test prompt, got %s", cfg.ExtraPrompt) 149 | } 150 | }) 151 | 152 | t.Run("config snapshot", func(t *testing.T) { 153 | tmpDir := t.TempDir() 154 | if err := os.Setenv("XDG_CONFIG_HOME", tmpDir); err != nil { 155 | t.Fatalf("failed to set XDG_CONFIG_HOME: %v", err) 156 | } 157 | defer func() { 158 | if err := os.Unsetenv("XDG_CONFIG_HOME"); err != nil { 159 | t.Errorf("failed to unset XDG_CONFIG_HOME: %v", err) 160 | } 161 | }() 162 | 163 | // Create test config directory 164 | configDir := filepath.Join(tmpDir, "llmscript") 165 | if err := os.MkdirAll(configDir, 0755); err != nil { 166 | t.Fatalf("failed to create config directory: %v", err) 167 | } 168 | 169 | cfg := DefaultConfig() 170 | cfg.LLM.Provider = "ollama" 171 | cfg.LLM.Ollama.Model = "llama2" 172 | cfg.LLM.Ollama.Host = "http://localhost:11434" 173 | cfg.LLM.Claude.APIKey = "" 174 | cfg.LLM.Claude.Model = "" 175 | cfg.LLM.OpenAI.APIKey = "" 176 | cfg.LLM.OpenAI.Model = "" 177 | cfg.MaxFixes = 5 178 | cfg.MaxAttempts = 2 179 | cfg.Timeout = 15 * time.Second 180 | cfg.ExtraPrompt = "Test prompt" 181 | 182 | if err := WriteConfig(cfg); err != nil { 183 | t.Fatalf("failed to write config: %v", err) 184 | } 185 | 186 | // Read the written file 187 | configPath := filepath.Join(configDir, "config.yaml") 188 | written, err := os.ReadFile(configPath) 189 | if err != nil { 190 | t.Fatalf("failed to read written config: %v", err) 191 | } 192 | 193 | // Compare with expected snapshot 194 | expected := `llm: 195 | provider: ollama 196 | ollama: 197 | model: llama2 198 | host: http://localhost:11434 199 | claude: 200 | api_key: "" 201 | model: "" 202 | openai: 203 | api_key: "" 204 | model: "" 205 | max_fixes: 5 206 | max_attempts: 2 207 | timeout: 15s 208 | additional_prompt: Test prompt 209 | ` 210 | if string(written) != expected { 211 | t.Errorf("config snapshot mismatch:\nExpected:\n%s\nGot:\n%s", expected, string(written)) 212 | } 213 | }) 214 | } 215 | --------------------------------------------------------------------------------