mirror of
https://github.com/lukaszraczylo/mcp-captcharesolver.git
synced 2026-10-12 00:45:28 +00:00
feat(llm): implement native gemini generateContent provider (vision + audio)
- Replace stub Vision/Transcribe with real REST implementation
- Endpoint: POST {base}/v1beta/models/{model}:generateContent
- Auth via x-goog-api-key header
- Vision: text part + inline_data per image (mime_type from Image.MediaType)
- Transcribe: transcription prompt + inline_data audio (default audio/wav)
- Model selection: opts.Model > audioModel (transcribe) > c.model
- maxTokens mapped to generationConfig.maxOutputTokens when non-zero
- Add httptest coverage: TestGeminiVision, TestGeminiTranscribe
- Re-enable gemini in tools.go audio capability check
This commit is contained in:
+100
-4
@@ -1,7 +1,12 @@
|
|||||||
package llm
|
package llm
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -14,10 +19,101 @@ type geminiClient struct {
|
|||||||
maxTokens int
|
maxTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *geminiClient) Vision(_ context.Context, _ string, _ []Image, _ Options) (string, error) {
|
// generateContent calls POST {base}/v1beta/models/{model}:generateContent.
|
||||||
return "", ErrUnsupported{Provider: "gemini"}
|
func (c *geminiClient) generateContent(ctx context.Context, model string, parts []map[string]any) (string, error) {
|
||||||
|
body := map[string]any{
|
||||||
|
"contents": []map[string]any{
|
||||||
|
{"parts": parts},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
if c.maxTokens > 0 {
|
||||||
|
body["generationConfig"] = map[string]any{"maxOutputTokens": c.maxTokens}
|
||||||
|
}
|
||||||
|
|
||||||
|
b, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
|
||||||
|
url := fmt.Sprintf("%s/v1beta/models/%s:generateContent", c.base, model)
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
req.Header.Set("x-goog-api-key", c.key)
|
||||||
|
|
||||||
|
resp, err := c.http.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
defer func() { _ = resp.Body.Close() }()
|
||||||
|
|
||||||
|
raw, err := io.ReadAll(resp.Body)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if resp.StatusCode >= 300 {
|
||||||
|
return "", fmt.Errorf("gemini HTTP %d: %s", resp.StatusCode, string(raw))
|
||||||
|
}
|
||||||
|
|
||||||
|
var out struct {
|
||||||
|
Candidates []struct {
|
||||||
|
Content struct {
|
||||||
|
Parts []struct {
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"parts"`
|
||||||
|
} `json:"content"`
|
||||||
|
} `json:"candidates"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
if len(out.Candidates) == 0 || len(out.Candidates[0].Content.Parts) == 0 {
|
||||||
|
return "", fmt.Errorf("gemini: empty response")
|
||||||
|
}
|
||||||
|
return out.Candidates[0].Content.Parts[0].Text, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *geminiClient) Transcribe(_ context.Context, _ []byte, _ string, _ Options) (string, error) {
|
func (c *geminiClient) Vision(ctx context.Context, prompt string, images []Image, opts Options) (string, error) {
|
||||||
return "", ErrUnsupported{Provider: "gemini"}
|
model := opts.Model
|
||||||
|
if model == "" {
|
||||||
|
model = c.model
|
||||||
|
}
|
||||||
|
|
||||||
|
parts := []map[string]any{
|
||||||
|
{"text": prompt},
|
||||||
|
}
|
||||||
|
for _, img := range images {
|
||||||
|
parts = append(parts, map[string]any{
|
||||||
|
"inline_data": map[string]any{
|
||||||
|
"mime_type": img.MediaType,
|
||||||
|
"data": base64.StdEncoding.EncodeToString(img.Data),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return c.generateContent(ctx, model, parts)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (c *geminiClient) Transcribe(ctx context.Context, audio []byte, _ string, opts Options) (string, error) {
|
||||||
|
model := opts.Model
|
||||||
|
if model == "" {
|
||||||
|
model = c.audioModel
|
||||||
|
}
|
||||||
|
if model == "" {
|
||||||
|
model = c.model
|
||||||
|
}
|
||||||
|
|
||||||
|
parts := []map[string]any{
|
||||||
|
{"text": "Transcribe the spoken characters in this audio captcha. Output only the transcription, nothing else."},
|
||||||
|
{
|
||||||
|
"inline_data": map[string]any{
|
||||||
|
"mime_type": "audio/wav",
|
||||||
|
"data": base64.StdEncoding.EncodeToString(audio),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
return c.generateContent(ctx, model, parts)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,94 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func geminiResponse(text string) []byte {
|
||||||
|
b, _ := json.Marshal(map[string]any{
|
||||||
|
"candidates": []map[string]any{
|
||||||
|
{"content": map[string]any{
|
||||||
|
"parts": []map[string]any{
|
||||||
|
{"text": text},
|
||||||
|
},
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGeminiVision(t *testing.T) {
|
||||||
|
const testKey = "test-api-key-vision"
|
||||||
|
imgData := []byte("fakeimagedata")
|
||||||
|
imgB64 := base64.StdEncoding.EncodeToString(imgData)
|
||||||
|
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
if !strings.Contains(r.URL.Path, ":generateContent") {
|
||||||
|
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||||
|
}
|
||||||
|
if got := r.Header.Get("x-goog-api-key"); got != testKey {
|
||||||
|
t.Errorf("x-goog-api-key = %q, want %q", got, testKey)
|
||||||
|
}
|
||||||
|
var body map[string]any
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||||
|
t.Fatalf("decode body: %v", err)
|
||||||
|
}
|
||||||
|
raw, _ := json.Marshal(body)
|
||||||
|
s := string(raw)
|
||||||
|
if !strings.Contains(s, "inline_data") {
|
||||||
|
t.Errorf("body missing inline_data: %s", s)
|
||||||
|
}
|
||||||
|
if !strings.Contains(s, imgB64) {
|
||||||
|
t.Errorf("body missing image base64: %s", s)
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write(geminiResponse("motorcycle"))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
c := &geminiClient{
|
||||||
|
http: srv.Client(),
|
||||||
|
base: srv.URL,
|
||||||
|
key: testKey,
|
||||||
|
model: "gemini-pro-vision",
|
||||||
|
}
|
||||||
|
|
||||||
|
got, err := c.Vision(context.Background(), "what is this?", []Image{
|
||||||
|
{MediaType: "image/png", Data: imgData},
|
||||||
|
}, Options{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Vision error: %v", err)
|
||||||
|
}
|
||||||
|
if got != "motorcycle" {
|
||||||
|
t.Errorf("Vision = %q, want %q", got, "motorcycle")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGeminiTranscribe(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write(geminiResponse("forty two"))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
c := &geminiClient{
|
||||||
|
http: srv.Client(),
|
||||||
|
base: srv.URL,
|
||||||
|
key: "test-key",
|
||||||
|
model: "gemini-pro",
|
||||||
|
}
|
||||||
|
|
||||||
|
got, err := c.Transcribe(context.Background(), []byte("audiodata"), "audio.wav", Options{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Transcribe error: %v", err)
|
||||||
|
}
|
||||||
|
if got != "forty two" {
|
||||||
|
t.Errorf("Transcribe = %q, want %q", got, "forty two")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -174,8 +174,7 @@ type SolverInfoOut struct {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (h *Handlers) solverInfo(_ context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, SolverInfoOut, error) {
|
func (h *Handlers) solverInfo(_ context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, SolverInfoOut, error) {
|
||||||
// gemini is a stub; re-add when implemented
|
audio := h.provider == "openai" || h.provider == "openai-compatible" || h.provider == "gemini"
|
||||||
audio := h.provider == "openai" || h.provider == "openai-compatible"
|
|
||||||
return nil, SolverInfoOut{
|
return nil, SolverInfoOut{
|
||||||
Provider: h.provider, Model: h.model,
|
Provider: h.provider, Model: h.model,
|
||||||
Capabilities: map[string]bool{
|
Capabilities: map[string]bool{
|
||||||
|
|||||||
Reference in New Issue
Block a user