subtitle-merge/gemini/main.go

113 lines
2.7 KiB
Go

package main
import (
"context"
"encoding/json"
"fmt"
"os"
"github.com/google/generative-ai-go/genai"
"google.golang.org/api/option"
)
// Response represents the JSON response structure
type Response struct {
Success bool `json:"success"`
Text string `json:"text,omitempty"`
Error string `json:"error,omitempty"`
}
// logToStderr logs messages to stderr for debugging
func logToStderr(format string, a ...interface{}) {
fmt.Fprintf(os.Stderr, format+"\n", a...)
}
func main() {
if len(os.Args) < 3 {
logToStderr("Usage: %s <api_key> <prompt> [image_path]", os.Args[0])
os.Exit(1)
}
apiKey := os.Args[1]
prompt := os.Args[2]
var imagePath string
if len(os.Args) > 3 {
imagePath = os.Args[3]
}
logToStderr("\n=== Gemini API Request ===")
logToStderr("API Key: %s...%s", apiKey[:4], apiKey[len(apiKey)-4:])
logToStderr("Prompt:\n%s", prompt)
ctx := context.Background()
client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
if err != nil {
handleError(fmt.Errorf("failed to create client: %v", err))
return
}
defer client.Close()
model := client.GenerativeModel("gemini-2.0-flash-exp")
var response *genai.GenerateContentResponse
if imagePath != "" {
logToStderr("\nProcessing image: %s", imagePath)
imageData, err := os.ReadFile(imagePath)
if err != nil {
handleError(fmt.Errorf("failed to read image: %v", err))
return
}
logToStderr("Image size: %d bytes", len(imageData))
model = client.GenerativeModel("gemini-2.0-flash-exp")
response, err = model.GenerateContent(ctx, genai.Text(prompt), genai.ImageData("image/jpeg", imageData))
} else {
logToStderr("\nProcessing text-only request")
response, err = model.GenerateContent(ctx, genai.Text(prompt))
}
if err != nil {
handleError(fmt.Errorf("failed to generate content: %v", err))
return
}
if len(response.Candidates) == 0 {
handleError(fmt.Errorf("no response candidates received"))
return
}
text := fmt.Sprintf("%v", response.Candidates[0].Content.Parts[0])
logToStderr("\n=== Gemini API Response ===")
logToStderr("Response:\n%s", text)
result := Response{
Success: true,
Text: text,
}
logToStderr("\n=== JSON Response ===")
prettyJSON, _ := json.MarshalIndent(result, "", " ")
logToStderr(string(prettyJSON))
// Output the final JSON to stdout without indentation
finalJSON, _ := json.Marshal(result)
fmt.Print(string(finalJSON))
}
func handleError(err error) {
logToStderr("\n=== Gemini API Error ===")
logToStderr("Error: %v", err)
errorResult := Response{
Success: false,
Error: err.Error(),
}
logToStderr("\n=== JSON Error Response ===")
prettyJSON, _ := json.MarshalIndent(errorResult, "", " ")
logToStderr(string(prettyJSON))
finalJSON, _ := json.Marshal(errorResult)
fmt.Print(string(finalJSON))
}