easyai-ai-gateway/apps/api/internal/clients/gemini.go

package clients

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type GeminiClient struct {
	HTTPClient *http.Client
}

func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error) {
	apiKey := credential(request.Candidate.Credentials, "apiKey", "api_key", "key", "token")
	if apiKey == "" {
		return Response{}, &ClientError{Code: "missing_credentials", Message: "gemini api key is required", Retryable: false}
	}
	body := geminiBody(request)
	raw, _ := json.Marshal(body)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, geminiURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey), bytes.NewReader(raw))
	if err != nil {
		return Response{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
	if err != nil {
		return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
	}
	responseStartedAt := time.Now()
	requestID := requestIDFromHTTPResponse(resp)
	result, err := decodeHTTPResponse(resp)
	responseFinishedAt := time.Now()
	if err != nil {
		return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
	}
	output := geminiResult(request, result)
	if requestID == "" {
		requestID = requestIDFromResult(output)
	}
	return Response{
		Result:             output,
		RequestID:          requestID,
		Usage:              geminiUsage(result),
		Progress:           providerProgress(request),
		ResponseStartedAt:  responseStartedAt,
		ResponseFinishedAt: responseFinishedAt,
		ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
	}, nil
}

func geminiURL(baseURL string, model string, apiKey string) string {
	base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if base == "" {
		base = "https://generativelanguage.googleapis.com"
	}
	if strings.HasSuffix(base, "/v1beta") {
		base = strings.TrimSuffix(base, "/v1beta")
	}
	escapedModel := url.PathEscape(model)
	return fmt.Sprintf("%s/v1beta/models/%s:generateContent?key=%s", base, escapedModel, url.QueryEscape(apiKey))
}

func geminiBody(request Request) map[string]any {
	if contents, ok := request.Body["contents"]; ok {
		return map[string]any{"contents": contents}
	}
	prompt := firstNonEmptyPrompt(request.Body, "")
	if prompt == "" {
		prompt = textFromMessages(request.Body)
	}
	return map[string]any{
		"contents": []any{map[string]any{
			"role":  "user",
			"parts": []any{map[string]any{"text": prompt}},
		}},
	}
}

func geminiResult(request Request, raw map[string]any) map[string]any {
	if request.ModelType == "image" {
		data := geminiImageData(raw)
		if len(data) == 0 {
			data = []any{map[string]any{"url": "/static/provider/gemini-image-placeholder.png"}}
		}
		return map[string]any{
			"id":      "gemini-image",
			"created": nowUnix(),
			"model":   request.Model,
			"data":    data,
			"raw":     raw,
		}
	}
	content := geminiText(raw)
	return map[string]any{
		"id":      "gemini-chat",
		"object":  "chat.completion",
		"created": nowUnix(),
		"model":   request.Model,
		"choices": []any{map[string]any{
			"index":         0,
			"finish_reason": "stop",
			"message":       map[string]any{"role": "assistant", "content": content},
		}},
		"usage": geminiUsageMap(raw),
		"raw":   raw,
	}
}

func textFromMessages(body map[string]any) string {
	messages, _ := body["messages"].([]any)
	parts := make([]string, 0, len(messages))
	for _, message := range messages {
		item, _ := message.(map[string]any)
		content := item["content"]
		switch typed := content.(type) {
		case string:
			parts = append(parts, typed)
		case []any:
			for _, part := range typed {
				partMap, _ := part.(map[string]any)
				if text, ok := partMap["text"].(string); ok {
					parts = append(parts, text)
				}
			}
		}
	}
	return strings.TrimSpace(strings.Join(parts, "\n"))
}

func geminiText(raw map[string]any) string {
	candidates, _ := raw["candidates"].([]any)
	for _, candidate := range candidates {
		candidateMap, _ := candidate.(map[string]any)
		content, _ := candidateMap["content"].(map[string]any)
		parts, _ := content["parts"].([]any)
		for _, part := range parts {
			partMap, _ := part.(map[string]any)
			if text, ok := partMap["text"].(string); ok && text != "" {
				return text
			}
		}
	}
	return ""
}

func geminiImageData(raw map[string]any) []any {
	candidates, _ := raw["candidates"].([]any)
	out := []any{}
	for _, candidate := range candidates {
		candidateMap, _ := candidate.(map[string]any)
		content, _ := candidateMap["content"].(map[string]any)
		parts, _ := content["parts"].([]any)
		for _, part := range parts {
			partMap, _ := part.(map[string]any)
			inline, _ := partMap["inlineData"].(map[string]any)
			if inline == nil {
				inline, _ = partMap["inline_data"].(map[string]any)
			}
			if data, ok := inline["data"].(string); ok && data != "" {
				out = append(out, map[string]any{"b64_json": data, "mime_type": inline["mimeType"]})
			}
		}
	}
	return out
}

func geminiUsage(raw map[string]any) Usage {
	usageMap := geminiUsageMap(raw)
	input := intFromAny(usageMap["prompt_tokens"])
	output := intFromAny(usageMap["completion_tokens"])
	total := intFromAny(usageMap["total_tokens"])
	return Usage{InputTokens: input, OutputTokens: output, TotalTokens: total}
}

func geminiUsageMap(raw map[string]any) map[string]any {
	meta, _ := raw["usageMetadata"].(map[string]any)
	input := intFromAny(meta["promptTokenCount"])
	output := intFromAny(meta["candidatesTokenCount"])
	total := intFromAny(meta["totalTokenCount"])
	if total == 0 {
		total = input + output
	}
	return map[string]any{"prompt_tokens": input, "completion_tokens": output, "total_tokens": total}
}