feat: implement AI gateway phase one runtime

2026-05-09 21:18:32 +08:00
parent a5e66e79cd
commit fdcdcd477b
46 changed files with 5678 additions and 768 deletions
@@ -0,0 +1,170 @@
+package clients
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/url"
+	"strings"
+)
+
+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, request.Candidate.ModelName, apiKey), bytes.NewReader(raw))
+	if err != nil {
+		return Response{}, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	resp, err := httpClient(c.HTTPClient).Do(req)
+	if err != nil {
+		return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
+	}
+	result, err := decodeHTTPResponse(resp)
+	if err != nil {
+		return Response{}, err
+	}
+	return Response{Result: geminiResult(request, result), Usage: geminiUsage(result), Progress: providerProgress(request)}, nil
+}
+
+func geminiURL(baseURL string, model string, apiKey string) string {
+	base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
+	if base == "" {
+		base = "https://generativelanguage.googleapis.com"
+	}
+	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}
+}