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
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
@@ -14,10 +19,101 @@ type geminiClient struct {
|
||||
maxTokens int
|
||||
}
|
||||
|
||||
func (c *geminiClient) Vision(_ context.Context, _ string, _ []Image, _ Options) (string, error) {
|
||||
return "", ErrUnsupported{Provider: "gemini"}
|
||||
// generateContent calls POST {base}/v1beta/models/{model}:generateContent.
|
||||
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) {
|
||||
return "", ErrUnsupported{Provider: "gemini"}
|
||||
func (c *geminiClient) Vision(ctx context.Context, prompt string, images []Image, opts Options) (string, error) {
|
||||
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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user