internal/codeimage: extract code from images

The codeimage package attempts to extract code from images using an
LLM.

The LLM's output is checked for syntactic validity, but it may not
compile. We can't actually compile it without a sandbox.

Perhaps in a followup we can typecheck it with go/types.
I'm not sure if we can use go/packages because it may invoke the compiler
for some configurations.

For #67.


Change-Id: Ibb9b65a6bdf7fc2cefa0a6eb341939b4a7456c89
Reviewed-on: https://go-review.googlesource.com/c/oscar/+/634295
Reviewed-by: Zvonimir Pavlinovic <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Hyang-Ah Hana Kim <[email protected]>

internal/codeimage/codeimage.go

@@ -0,0 +1,76 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package codeimage extracts Go code from images.
+package codeimage
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"go/format"
+	"io"
+	"net/http"
+	"strconv"
+
+	"golang.org/x/oscar/internal/llm"
+)
+
+// InURL looks for code in the image referred to by url.
+// It fetches the url using the given HTTP client, then calls [InBlob].
+func InURL(ctx context.Context, url string, client *http.Client, cgen llm.ContentGenerator) (string, error) {
+	res, err := client.Get(url)
+	if err != nil {
+		return "", err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusOK {
+		return "", fmt.Errorf("%s returned status %s", url, res.Status)
+	}
+	data, err := io.ReadAll(res.Body)
+	if err != nil {
+		return "", err
+	}
+	return InBlob(ctx, llm.Blob{
+		MIMEType: res.Header.Get("Content-Type"),
+		Data:     data,
+	}, cgen)
+}
+
+// InBlob looks for code in the given blob.
+// On success, it returns a properly formatted Go program or program fragment.
+func InBlob(ctx context.Context, blob llm.Blob, cgen llm.ContentGenerator) (string, error) {
+	parts := []llm.Part{
+		llm.Text(instructions),
+		blob,
+	}
+	schema := &llm.Schema{
+		Type:        llm.TypeString,
+		Description: "the program",
+	}
+	// Retry several times in the hope that the LLM will eventually produce a valid program.
+	for try := 0; try < 3; try++ {
+		output, err := cgen.GenerateContent(ctx, schema, parts)
+		if err != nil {
+			return "", err
+		}
+		unq, err := strconv.Unquote(output)
+		if err != nil {
+			return "", fmt.Errorf("strconv.Unquote: %w", err)
+		}
+		fbytes, err := format.Source([]byte(unq))
+		if err != nil {
+			// Retry if it isn't a valid program.
+			continue
+		}
+		return string(fbytes), nil
+	}
+	return "", errors.New("could not produce a valid Go program")
+}
+
+const instructions = `
+The following image contains code in the Go programming language.
+Extract the Go code from the image.
+Make sure you produce a syntactically valid Go program.
+`

internal/codeimage/codeimage_test.go

@@ -0,0 +1,65 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package codeimage
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"slices"
+	"strconv"
+	"strings"
+	"testing"
+
+	"golang.org/x/oscar/internal/llm"
+)
+
+var ctx = context.Background()
+
+func TestInBlob(t *testing.T) {
+	imageFiles, err := filepath.Glob(filepath.Join("testdata", "*.png"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, ifile := range imageFiles {
+		noext := strings.TrimSuffix(ifile, ".png")
+		t.Run(filepath.Base(noext), func(t *testing.T) {
+			wbytes, err := os.ReadFile(noext + ".go")
+			if err != nil {
+				t.Fatal(err)
+			}
+			gen := newMockGenerator(string(wbytes))
+			data, err := os.ReadFile(ifile)
+			if err != nil {
+				t.Fatal(err)
+			}
+			blob := llm.Blob{MIMEType: "image/png", Data: data}
+			got, err := InBlob(ctx, blob, gen)
+			if err != nil {
+				t.Fatal(err)
+			}
+			got = removeBlankLines(got)
+			want := removeBlankLines(string(wbytes))
+			if got != want {
+				t.Errorf("\ngot:\n%s\nwant:\n%s", got, want)
+			}
+		})
+	}
+}
+
+func removeBlankLines(s string) string {
+	lines := strings.Split(s, "\n")
+	lines = slices.DeleteFunc(lines, func(s string) bool {
+		return len(strings.TrimSpace(s)) == 0
+	})
+	return strings.Join(lines, "\n")
+}
+
+func newMockGenerator(result string) llm.ContentGenerator {
+	return llm.TestContentGenerator("imageMock",
+		func(context.Context, *llm.Schema, []llm.Part) (string, error) {
+			return strconv.Quote(result), nil
+		})
+}

internal/codeimage/testdata/constants.go

@@ -0,0 +1,25 @@
+package main
+
+import "fmt"
+
+func needFloat(x float64) float64 {
+	return x * 0.1
+}
+
+func needInt(x int) int {
+	return x*10 + 1
+}
+
+func main() {
+	fmt.Println(needInt(Small))
+	fmt.Println(needFloat(Small))
+	fmt.Println(needFloat(Big))
+}
+
+const (
+	// A big binary number.
+	Big = 1 << 100
+
+	// A small one.
+	Small = Big >> 99
+)

internal/codeimage/testdata/hello-world.go

@@ -0,0 +1,7 @@
+package main
+
+import "fmt"
+
+func main() {
+	fmt.Println("Hello, world.")
+}

internal/codeimage/testdata/issue-52950.go

@@ -0,0 +1,16 @@
+package main
+
+import "fmt"
+
+func main() {
+	var pi float32 = 3.14
+	switch {
+	case pi == 3.14:
+		fmt.Println("pi == 3.14")
+		fallthrough
+	case pi == 11.00:
+		fmt.Println("pi == 11.00")
+	default:
+		fmt.Println("i'm here")
+	}
+}