diff --git a/internal/captcha/grid.go b/internal/captcha/grid.go index 79695d2..089d5e1 100644 --- a/internal/captcha/grid.go +++ b/internal/captcha/grid.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "sort" "github.com/lukaszraczylo/captcha-solver-mcp/internal/imageutil" "github.com/lukaszraczylo/captcha-solver-mcp/internal/llm" @@ -17,6 +18,9 @@ type GridOptions struct { Cols int ImageWidth int ImageHeight int + Samples int + Upscale int + Annotate bool } // GridResult lists tiles to click (row-major 0-based) plus optional pixel @@ -25,17 +29,20 @@ type GridResult struct { Tiles []int `json:"tiles"` Centroids []imageutil.Point `json:"centroids,omitempty"` Confidence float64 `json:"confidence"` + NoTargets bool `json:"no_targets"` RecheckRecommended bool `json:"recheck_recommended"` } func gridPrompt(o GridOptions, strict bool) string { p := fmt.Sprintf( "You are solving an image-grid captcha. The image is a %dx%d grid of tiles, "+ - "numbered row-major starting at 0 (top-left). Instruction: %q. "+ - "Return the 0-based indices of every tile that matches the instruction. "+ - "If after selection more matching tiles would likely appear (dynamic grid), set recheck_recommended true. "+ - `Respond ONLY with compact JSON: {"tiles":[..],"confidence":<0..1>,"recheck_recommended":}.`, - o.Rows, o.Cols, o.Instruction) + "numbered row-major starting at 0 (top-left) through %d (numbers may be overlaid on the image). "+ + "Instruction: %q. "+ + "Select EVERY cell that contains ANY part of the target, including small or partial portions. "+ + "If no cell contains the target, return an empty tiles array and set no_targets true. "+ + "If this is a dynamic grid where new matching tiles may appear after selection, set recheck_recommended true. ", + o.Rows, o.Cols, o.Rows*o.Cols-1, o.Instruction) + p += `Respond ONLY with compact JSON: {"tiles":[..],"confidence":<0..1>,"no_targets":,"recheck_recommended":}.` if strict { p += " You MUST output valid JSON only. No prose." } @@ -43,23 +50,118 @@ func gridPrompt(o GridOptions, strict bool) string { } // SolveGrid asks the vision model which tiles to click and (when grid geometry -// is supplied) maps them to pixel centroids for humanized clicking. +// is supplied) maps them to pixel centroids for humanized clicking. When +// o.Samples > 1 it self-consistency votes across N solves, keeping tiles that +// appear in a strict majority of successful samples. func SolveGrid(ctx context.Context, c llm.Client, img []byte, mediaType string, o GridOptions) (GridResult, error) { + img, mediaType = preprocessGrid(img, mediaType, o) images := []llm.Image{{MediaType: mediaType, Data: img}} + + if o.Samples <= 1 { + res, err := solveGridOnce(ctx, c, images, o, llm.Options{}) + if err != nil { + return GridResult{}, err + } + res.Centroids = mapCentroids(res.Tiles, o) + return res, nil + } + + opts := llm.Options{Temperature: 0.5} + results := make([]GridResult, 0, o.Samples) + for i := 0; i < o.Samples; i++ { + res, err := solveGridOnce(ctx, c, images, o, opts) + if err != nil { + continue + } + results = append(results, res) + } + if len(results) == 0 { + return GridResult{}, fmt.Errorf("model did not return parseable tile selection") + } + voted := voteGrid(results) + voted.Centroids = mapCentroids(voted.Tiles, o) + return voted, nil +} + +// preprocessGrid optionally upscales then annotates the screenshot. On any +// preprocessing error it falls back to the original bytes so a solve is never +// failed by an optional enhancement. +func preprocessGrid(img []byte, mediaType string, o GridOptions) ([]byte, string) { + if o.Upscale > 1 { + if up, mt, err := imageutil.Upscale(img, o.Upscale); err == nil { + img, mediaType = up, mt + } + } + if o.Annotate { + if ann, err := imageutil.AnnotateGrid(img, o.Rows, o.Cols); err == nil { + img, mediaType = ann, "image/png" + } + } + return img, mediaType +} + +// solveGridOnce performs a single solve with one stricter retry on parse fail. +func solveGridOnce(ctx context.Context, c llm.Client, images []llm.Image, o GridOptions, opts llm.Options) (GridResult, error) { for attempt := 0; attempt < 2; attempt++ { - raw, err := c.Vision(ctx, gridPrompt(o, attempt > 0), images, llm.Options{}) + raw, err := c.Vision(ctx, gridPrompt(o, attempt > 0), images, opts) if err != nil { return GridResult{}, err } var res GridResult if err := json.Unmarshal([]byte(extractJSON(raw)), &res); err == nil && res.Tiles != nil { - res.Centroids = mapCentroids(res.Tiles, o) return res, nil } } return GridResult{}, fmt.Errorf("model did not return parseable tile selection") } +// voteGrid keeps tiles appearing in a strict majority of successful samples. +// Confidence is the average agreement of the chosen tiles; NoTargets is true +// when a majority of samples reported no targets. +func voteGrid(results []GridResult) GridResult { + n := len(results) + tileCounts := make(map[int]int) + noTargets := 0 + recheck := 0 + for _, r := range results { + seen := make(map[int]struct{}, len(r.Tiles)) + for _, t := range r.Tiles { + if _, dup := seen[t]; dup { + continue + } + seen[t] = struct{}{} + tileCounts[t]++ + } + if r.NoTargets { + noTargets++ + } + if r.RecheckRecommended { + recheck++ + } + } + + tiles := make([]int, 0, len(tileCounts)) + var agreeSum float64 + for t, cnt := range tileCounts { + if cnt*2 > n { // strict majority: more than half + tiles = append(tiles, t) + agreeSum += float64(cnt) / float64(n) + } + } + sort.Ints(tiles) + + var conf float64 + if len(tiles) > 0 { + conf = agreeSum / float64(len(tiles)) + } + return GridResult{ + Tiles: tiles, + Confidence: conf, + NoTargets: noTargets*2 > n, + RecheckRecommended: recheck*2 > n, + } +} + func mapCentroids(tiles []int, o GridOptions) []imageutil.Point { if o.Rows <= 0 || o.Cols <= 0 || o.ImageWidth <= 0 || o.ImageHeight <= 0 { return nil diff --git a/internal/captcha/grid_test.go b/internal/captcha/grid_test.go index d8693a5..b4fddce 100644 --- a/internal/captcha/grid_test.go +++ b/internal/captcha/grid_test.go @@ -1,12 +1,32 @@ package captcha import ( + "bytes" "context" + "image" + "image/color" + "image/png" "testing" "github.com/lukaszraczylo/captcha-solver-mcp/internal/llm" ) +// smallPNG returns a tiny valid PNG for exercising preprocessing paths. +func smallPNG(t *testing.T) []byte { + t.Helper() + img := image.NewRGBA(image.Rect(0, 0, 6, 6)) + for y := 0; y < 6; y++ { + for x := 0; x < 6; x++ { + img.Set(x, y, color.RGBA{R: uint8(x * 40), G: uint8(y * 40), B: 80, A: 255}) + } + } + var buf bytes.Buffer + if err := png.Encode(&buf, img); err != nil { + t.Fatalf("encode png: %v", err) + } + return buf.Bytes() +} + func TestSolveGrid(t *testing.T) { f := &llm.Fake{VisionResp: `{"tiles":[0,4,8],"confidence":0.8,"recheck_recommended":true}`} res, err := SolveGrid(context.Background(), f, []byte{1}, "image/png", GridOptions{ @@ -25,3 +45,53 @@ func TestSolveGrid(t *testing.T) { t.Fatal("recheck flag lost") } } + +func TestSolveGridVoting(t *testing.T) { + f := &llm.Fake{VisionResponses: []string{ + `{"tiles":[0,1],"confidence":0.9}`, + `{"tiles":[0,2],"confidence":0.9}`, + `{"tiles":[0,1],"confidence":0.9}`, + }} + res, err := SolveGrid(context.Background(), f, []byte{1}, "image/png", GridOptions{ + Instruction: "select all buses", Rows: 3, Cols: 3, Samples: 3, + }) + if err != nil { + t.Fatal(err) + } + // 0 in 3/3, 1 in 2/3 (>half), 2 in 1/3 (dropped). + if len(res.Tiles) != 2 || res.Tiles[0] != 0 || res.Tiles[1] != 1 { + t.Fatalf("voted tiles %+v, want [0 1]", res.Tiles) + } + if len(f.VisionCalls) != 3 { + t.Fatalf("expected 3 vision calls, got %d", len(f.VisionCalls)) + } +} + +func TestSolveGridNoTargets(t *testing.T) { + f := &llm.Fake{VisionResp: `{"tiles":[],"confidence":0.9,"no_targets":true}`} + res, err := SolveGrid(context.Background(), f, []byte{1}, "image/png", GridOptions{ + Instruction: "select all buses", Rows: 3, Cols: 3, + }) + if err != nil { + t.Fatal(err) + } + if !res.NoTargets { + t.Fatal("no_targets flag lost") + } + if len(res.Tiles) != 0 { + t.Fatalf("expected empty tiles, got %+v", res.Tiles) + } +} + +func TestSolveGridPreprocessing(t *testing.T) { + f := &llm.Fake{VisionResp: `{"tiles":[0,4],"confidence":0.8}`} + res, err := SolveGrid(context.Background(), f, smallPNG(t), "image/png", GridOptions{ + Instruction: "select all buses", Rows: 3, Cols: 3, Annotate: true, Upscale: 2, + }) + if err != nil { + t.Fatal(err) + } + if len(res.Tiles) != 2 { + t.Fatalf("expected 2 tiles, got %+v", res.Tiles) + } +} diff --git a/internal/captcha/rotation.go b/internal/captcha/rotation.go index 4db8b45..431fa2d 100644 --- a/internal/captcha/rotation.go +++ b/internal/captcha/rotation.go @@ -21,7 +21,8 @@ func SolveRotation(ctx context.Context, c llm.Client, img []byte, mediaType, ins "You are solving a rotation/choice puzzle captcha. Instruction: %q. "+ "If asked to pick one image, return its 0-based index as choice. "+ "If asked to rotate, return the number of clockwise steps as rotations. "+ - `Respond ONLY with compact JSON: {"choice":,"rotations":,"confidence":<0..1>}.`, + "First reason step-by-step about the correct orientation/choice, then give the final answer "+ + `as the last thing you output: compact JSON {"choice":,"rotations":,"confidence":<0..1>}.`, instruction) raw, err := c.Vision(ctx, prompt, []llm.Image{{MediaType: mediaType, Data: img}}, llm.Options{}) if err != nil { diff --git a/internal/captcha/text.go b/internal/captcha/text.go index de70d0c..f674db6 100644 --- a/internal/captcha/text.go +++ b/internal/captcha/text.go @@ -4,7 +4,9 @@ import ( "context" "encoding/json" "fmt" + "sort" + "github.com/lukaszraczylo/captcha-solver-mcp/internal/imageutil" "github.com/lukaszraczylo/captcha-solver-mcp/internal/llm" ) @@ -13,6 +15,8 @@ type TextOptions struct { Hint string Charset string Length int + Samples int + Upscale int CaseSensitive bool } @@ -36,6 +40,7 @@ func textPrompt(o TextOptions, strict bool) string { if o.Hint != "" { p += "Hint: " + o.Hint + ". " } + p += "Carefully distinguish visually similar glyphs: 0 vs O, 1 vs l vs I, 5 vs S, 2 vs Z, 8 vs B. " p += `Respond ONLY with compact JSON: {"text":"","confidence":<0..1>}.` if strict { p += " You MUST output valid JSON and nothing else. Do not explain." @@ -43,11 +48,40 @@ func textPrompt(o TextOptions, strict bool) string { return p } -// SolveText runs OCR via the vision model with one stricter retry on parse failure. +// SolveText runs OCR via the vision model with one stricter retry on parse +// failure. When o.Samples > 1 it self-consistency votes across N solves, +// returning the most frequent text (tie-break by highest confidence). func SolveText(ctx context.Context, c llm.Client, img []byte, mediaType string, o TextOptions) (TextResult, error) { + if o.Upscale > 1 { + if up, mt, err := imageutil.Upscale(img, o.Upscale); err == nil { + img, mediaType = up, mt + } + } images := []llm.Image{{MediaType: mediaType, Data: img}} + + if o.Samples <= 1 { + return solveTextOnce(ctx, c, images, o, llm.Options{}) + } + + opts := llm.Options{Temperature: 0.5} + results := make([]TextResult, 0, o.Samples) + for i := 0; i < o.Samples; i++ { + res, err := solveTextOnce(ctx, c, images, o, opts) + if err != nil { + continue + } + results = append(results, res) + } + if len(results) == 0 { + return TextResult{}, fmt.Errorf("model did not return parseable captcha text") + } + return voteText(results), nil +} + +// solveTextOnce performs a single solve with one stricter retry on parse fail. +func solveTextOnce(ctx context.Context, c llm.Client, images []llm.Image, o TextOptions, opts llm.Options) (TextResult, error) { for attempt := 0; attempt < 2; attempt++ { - raw, err := c.Vision(ctx, textPrompt(o, attempt > 0), images, llm.Options{}) + raw, err := c.Vision(ctx, textPrompt(o, attempt > 0), images, opts) if err != nil { return TextResult{}, err } @@ -58,3 +92,37 @@ func SolveText(ctx context.Context, c llm.Client, img []byte, mediaType string, } return TextResult{}, fmt.Errorf("model did not return parseable captcha text") } + +// voteText returns the most frequent text across results. Ties break by the +// highest confidence seen for the candidate; remaining ties break +// deterministically by text. Confidence is the fraction of results that agree +// with the winner. +func voteText(results []TextResult) TextResult { + counts := make(map[string]int) + bestConf := make(map[string]float64) + for _, r := range results { + counts[r.Text]++ + if r.Confidence > bestConf[r.Text] { + bestConf[r.Text] = r.Confidence + } + } + texts := make([]string, 0, len(counts)) + for t := range counts { + texts = append(texts, t) + } + sort.Slice(texts, func(i, j int) bool { + a, b := texts[i], texts[j] + if counts[a] != counts[b] { + return counts[a] > counts[b] + } + if bestConf[a] != bestConf[b] { + return bestConf[a] > bestConf[b] + } + return a < b + }) + winner := texts[0] + return TextResult{ + Text: winner, + Confidence: float64(counts[winner]) / float64(len(results)), + } +} diff --git a/internal/captcha/text_test.go b/internal/captcha/text_test.go index ad68658..1895737 100644 --- a/internal/captcha/text_test.go +++ b/internal/captcha/text_test.go @@ -31,3 +31,21 @@ func TestSolveTextRetriesOnGarbage(t *testing.T) { t.Fatalf("expected 2 attempts (one retry), got %d", len(f.VisionCalls)) } } + +func TestSolveTextVoting(t *testing.T) { + f := &llm.Fake{VisionResponses: []string{ + `{"text":"ab12","confidence":0.7}`, + `{"text":"ab12","confidence":0.7}`, + `{"text":"ab13","confidence":0.9}`, + }} + res, err := SolveText(context.Background(), f, []byte{1}, "image/png", TextOptions{Samples: 3}) + if err != nil { + t.Fatal(err) + } + if res.Text != "ab12" { + t.Fatalf("voted text %q, want ab12", res.Text) + } + if len(f.VisionCalls) != 3 { + t.Fatalf("expected 3 vision calls, got %d", len(f.VisionCalls)) + } +}