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,120 @@
+package clients
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+func TestOpenAIClientChatContract(t *testing.T) {
+	var gotPath string
+	var gotAuth string
+	var gotModel string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.Path
+		gotAuth = r.Header.Get("Authorization")
+		var body map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		gotModel, _ = body["model"].(string)
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"id":     "chatcmpl-test",
+			"object": "chat.completion",
+			"model":  gotModel,
+			"choices": []any{map[string]any{
+				"message": map[string]any{"role": "assistant", "content": "ok"},
+			}},
+			"usage": map[string]any{"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5},
+		})
+	}))
+	defer server.Close()
+
+	response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
+		Kind:  "chat.completions",
+		Model: "openai:gpt-4o-mini",
+		Body:  map[string]any{"model": "openai:gpt-4o-mini", "messages": []any{map[string]any{"role": "user", "content": "ping"}}},
+		Candidate: store.RuntimeModelCandidate{
+			BaseURL:     server.URL,
+			ModelName:   "gpt-4o-mini",
+			Credentials: map[string]any{"apiKey": "test-key"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("run openai client: %v", err)
+	}
+	if gotPath != "/chat/completions" || gotAuth != "Bearer test-key" || gotModel != "gpt-4o-mini" {
+		t.Fatalf("unexpected request path=%s auth=%s model=%s", gotPath, gotAuth, gotModel)
+	}
+	if response.Usage.TotalTokens != 5 || response.Result["id"] != "chatcmpl-test" {
+		t.Fatalf("unexpected response: %+v", response)
+	}
+}
+
+func TestGeminiClientChatContract(t *testing.T) {
+	var gotPath string
+	var gotKey string
+	var gotText string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotPath = r.URL.Path
+		gotKey = r.URL.Query().Get("key")
+		var body map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		contents, _ := body["contents"].([]any)
+		first, _ := contents[0].(map[string]any)
+		parts, _ := first["parts"].([]any)
+		part, _ := parts[0].(map[string]any)
+		gotText, _ = part["text"].(string)
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"candidates": []any{map[string]any{
+				"content": map[string]any{
+					"parts": []any{map[string]any{"text": "gemini ok"}},
+				},
+			}},
+			"usageMetadata": map[string]any{
+				"promptTokenCount":     4,
+				"candidatesTokenCount": 6,
+				"totalTokenCount":      10,
+			},
+		})
+	}))
+	defer server.Close()
+
+	response, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
+		Kind:  "chat.completions",
+		Model: "gemini:gemini-2.5-flash",
+		Body: map[string]any{
+			"model":    "gemini:gemini-2.5-flash",
+			"messages": []any{map[string]any{"role": "user", "content": "ping"}},
+		},
+		Candidate: store.RuntimeModelCandidate{
+			BaseURL:     server.URL,
+			ModelName:   "gemini-2.5-flash",
+			ModelType:   "chat",
+			Credentials: map[string]any{"apiKey": "gemini-key"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("run gemini client: %v", err)
+	}
+	if gotPath != "/v1beta/models/gemini-2.5-flash:generateContent" || gotKey != "gemini-key" || gotText != "ping" {
+		t.Fatalf("unexpected request path=%s key=%s text=%s", gotPath, gotKey, gotText)
+	}
+	if response.Usage.TotalTokens != 10 || extractText(response.Result) != "gemini ok" {
+		t.Fatalf("unexpected response: %+v", response)
+	}
+}
+
+func extractText(result map[string]any) string {
+	choices, _ := result["choices"].([]any)
+	choice, _ := choices[0].(map[string]any)
+	message, _ := choice["message"].(map[string]any)
+	text, _ := message["content"].(string)
+	return text
+}
@@ -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}
+}
@@ -0,0 +1,133 @@
+package clients
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"math"
+	"net/http"
+	"strings"
+	"time"
+)
+
+func credential(candidate map[string]any, keys ...string) string {
+	for _, key := range keys {
+		if value, ok := candidate[key].(string); ok && strings.TrimSpace(value) != "" {
+			return strings.TrimSpace(value)
+		}
+	}
+	return ""
+}
+
+func boolValue(body map[string]any, key string) bool {
+	value, _ := body[key].(bool)
+	return value
+}
+
+func stringValue(body map[string]any, key string) string {
+	value, _ := body[key].(string)
+	return strings.TrimSpace(value)
+}
+
+func intValue(body map[string]any, key string, fallback int) int {
+	switch value := body[key].(type) {
+	case float64:
+		return int(math.Round(value))
+	case int:
+		return value
+	default:
+		return fallback
+	}
+}
+
+func decodeHTTPResponse(resp *http.Response) (map[string]any, error) {
+	defer resp.Body.Close()
+	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return nil, &ClientError{
+			Code:       statusCodeName(resp.StatusCode),
+			Message:    errorMessage(raw, resp.Status),
+			StatusCode: resp.StatusCode,
+			Retryable:  HTTPRetryable(resp.StatusCode),
+		}
+	}
+	var out map[string]any
+	if len(raw) == 0 {
+		return map[string]any{}, nil
+	}
+	if err := json.Unmarshal(raw, &out); err != nil {
+		return nil, &ClientError{Code: "invalid_response", Message: err.Error(), Retryable: false}
+	}
+	return out, nil
+}
+
+func usageFromOpenAI(result map[string]any) Usage {
+	usage, _ := result["usage"].(map[string]any)
+	input := intFromAny(firstPresent(usage["prompt_tokens"], usage["input_tokens"]))
+	output := intFromAny(firstPresent(usage["completion_tokens"], usage["output_tokens"]))
+	total := intFromAny(usage["total_tokens"])
+	if total == 0 {
+		total = input + output
+	}
+	return Usage{InputTokens: input, OutputTokens: output, TotalTokens: total}
+}
+
+func intFromAny(value any) int {
+	switch typed := value.(type) {
+	case float64:
+		return int(math.Round(typed))
+	case int:
+		return typed
+	case int64:
+		return int(typed)
+	default:
+		return 0
+	}
+}
+
+func firstPresent(values ...any) any {
+	for _, value := range values {
+		if value != nil {
+			return value
+		}
+	}
+	return nil
+}
+
+func errorMessage(raw []byte, fallback string) string {
+	if len(raw) == 0 {
+		return fallback
+	}
+	var parsed map[string]any
+	if json.Unmarshal(raw, &parsed) == nil {
+		if errObj, ok := parsed["error"].(map[string]any); ok {
+			if message, ok := errObj["message"].(string); ok {
+				return message
+			}
+		}
+		if message, ok := parsed["message"].(string); ok {
+			return message
+		}
+	}
+	return string(raw)
+}
+
+func statusCodeName(status int) string {
+	switch status {
+	case http.StatusTooManyRequests:
+		return "rate_limit"
+	case http.StatusRequestTimeout:
+		return "timeout"
+	case http.StatusUnauthorized, http.StatusForbidden:
+		return "auth_failed"
+	default:
+		if status >= 500 {
+			return "server_error"
+		}
+		return fmt.Sprintf("http_%d", status)
+	}
+}
+
+func nowUnix() int64 {
+	return time.Now().Unix()
+}
@@ -0,0 +1,87 @@
+package clients
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"strings"
+)
+
+type OpenAIClient struct {
+	HTTPClient *http.Client
+}
+
+func (c OpenAIClient) 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: "openai api key is required", Retryable: false}
+	}
+	endpoint := openAIEndpoint(request.Kind)
+	if endpoint == "" {
+		return Response{}, &ClientError{Code: "unsupported_kind", Message: "unsupported openai request kind", Retryable: false}
+	}
+	body := cloneBody(request.Body)
+	body["model"] = request.Candidate.ModelName
+	raw, _ := json.Marshal(body)
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(request.Candidate.BaseURL, endpoint), bytes.NewReader(raw))
+	if err != nil {
+		return Response{}, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+apiKey)
+	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: result, Usage: usageFromOpenAI(result), Progress: providerProgress(request)}, nil
+}
+
+func openAIEndpoint(kind string) string {
+	switch kind {
+	case "chat.completions":
+		return "/chat/completions"
+	case "responses":
+		return "/responses"
+	case "images.generations":
+		return "/images/generations"
+	case "images.edits":
+		return "/images/edits"
+	default:
+		return ""
+	}
+}
+
+func cloneBody(body map[string]any) map[string]any {
+	out := map[string]any{}
+	for key, value := range body {
+		out[key] = value
+	}
+	return out
+}
+
+func joinURL(base string, path string) string {
+	base = strings.TrimRight(strings.TrimSpace(base), "/")
+	if base == "" {
+		base = "https://api.openai.com/v1"
+	}
+	return base + path
+}
+
+func httpClient(client *http.Client) *http.Client {
+	if client != nil {
+		return client
+	}
+	return http.DefaultClient
+}
+
+func providerProgress(request Request) []Progress {
+	return []Progress{
+		{Phase: "submitting", Progress: 0.35, Message: "provider request submitted", Payload: map[string]any{"clientId": request.Candidate.ClientID}},
+		{Phase: "fetching_result", Progress: 0.8, Message: "provider response received", Payload: map[string]any{"provider": request.Candidate.Provider}},
+	}
+}
@@ -0,0 +1,119 @@
+package clients
+
+import (
+	"context"
+	"fmt"
+	"strings"
+)
+
+type SimulationClient struct{}
+
+func (c SimulationClient) Run(ctx context.Context, request Request) (Response, error) {
+	_ = ctx
+	profile := simulationProfile(request)
+	if profile == "retryable_failure" {
+		return Response{}, &ClientError{Code: "server_error", Message: "simulated retryable failure", Retryable: true}
+	}
+	if profile == "fatal_failure" || profile == "non_retryable_failure" {
+		return Response{}, &ClientError{Code: "bad_request", Message: "simulated non-retryable failure", Retryable: false}
+	}
+	return Response{
+		Result:   simulatedResult(request),
+		Usage:    simulatedUsage(request),
+		Progress: simulatedProgress(request),
+	}, nil
+}
+
+func simulationProfile(request Request) string {
+	if value := stringValue(request.Candidate.Credentials, "simulationFailure"); value != "" {
+		return value
+	}
+	if value := stringValue(request.Candidate.PlatformConfig, "simulationFailure"); value != "" {
+		return value
+	}
+	if value := stringValue(request.Body, "simulationProfile"); value != "" {
+		return value
+	}
+	if value := stringValue(request.Body, "testProfile"); value != "" {
+		return value
+	}
+	return "success"
+}
+
+func simulatedResult(request Request) map[string]any {
+	switch request.Kind {
+	case "chat.completions":
+		return map[string]any{
+			"id":      "chatcmpl-simulated",
+			"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": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
+				},
+			}},
+			"usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20},
+		}
+	case "responses":
+		return map[string]any{
+			"id":          "resp-simulated",
+			"object":      "response",
+			"created_at":  nowUnix(),
+			"model":       request.Model,
+			"output_text": fmt.Sprintf("simulation response from %s", request.Candidate.Provider),
+			"usage":       map[string]any{"input_tokens": 12, "output_tokens": 8, "total_tokens": 20},
+		}
+	case "images.edits":
+		return map[string]any{
+			"id":      "img-edit-simulated",
+			"created": nowUnix(),
+			"model":   request.Model,
+			"data": []any{map[string]any{
+				"url":            "/static/simulation/image-edit.png",
+				"revised_prompt": firstNonEmptyPrompt(request.Body, "simulation image edit"),
+			}},
+		}
+	default:
+		return map[string]any{
+			"id":      "img-simulated",
+			"created": nowUnix(),
+			"model":   request.Model,
+			"data": []any{map[string]any{
+				"url":            "/static/simulation/image.png",
+				"revised_prompt": firstNonEmptyPrompt(request.Body, "simulation image"),
+			}},
+		}
+	}
+}
+
+func simulatedUsage(request Request) Usage {
+	if request.ModelType == "chat" || request.Kind == "responses" {
+		return Usage{InputTokens: 12, OutputTokens: 8, TotalTokens: 20}
+	}
+	return Usage{}
+}
+
+func simulatedProgress(request Request) []Progress {
+	provider := request.Candidate.Provider
+	if provider == "" {
+		provider = "simulation"
+	}
+	return []Progress{
+		{Phase: "normalizing", Progress: 0.2, Message: "request normalized", Payload: map[string]any{"provider": provider}},
+		{Phase: "submitting", Progress: 0.55, Message: "simulation client submitted", Payload: map[string]any{"clientId": request.Candidate.ClientID}},
+		{Phase: "fetching_result", Progress: 0.85, Message: "simulation result ready", Payload: map[string]any{"kind": request.Kind}},
+	}
+}
+
+func firstNonEmptyPrompt(body map[string]any, fallback string) string {
+	for _, key := range []string{"prompt", "input"} {
+		if value := strings.TrimSpace(stringValue(body, key)); value != "" {
+			return value
+		}
+	}
+	return fallback
+}
@@ -0,0 +1,72 @@
+package clients
+
+import (
+	"context"
+	"errors"
+	"net/http"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+type Request struct {
+	Kind      string
+	ModelType string
+	Model     string
+	Body      map[string]any
+	Candidate store.RuntimeModelCandidate
+	Stream    bool
+}
+
+type Response struct {
+	Result   map[string]any
+	Usage    Usage
+	Progress []Progress
+}
+
+type Usage struct {
+	InputTokens  int
+	OutputTokens int
+	TotalTokens  int
+}
+
+type Progress struct {
+	Phase    string
+	Progress float64
+	Message  string
+	Payload  map[string]any
+}
+
+type Client interface {
+	Run(ctx context.Context, request Request) (Response, error)
+}
+
+type ClientError struct {
+	Code       string
+	Message    string
+	StatusCode int
+	Retryable  bool
+}
+
+func (e *ClientError) Error() string {
+	if e.Message != "" {
+		return e.Message
+	}
+	return e.Code
+}
+
+func IsRetryable(err error) bool {
+	var clientErr *ClientError
+	return errors.As(err, &clientErr) && clientErr.Retryable
+}
+
+func ErrorCode(err error) string {
+	var clientErr *ClientError
+	if errors.As(err, &clientErr) && clientErr.Code != "" {
+		return clientErr.Code
+	}
+	return "client_error"
+}
+
+func HTTPRetryable(status int) bool {
+	return status == http.StatusTooManyRequests || status == http.StatusRequestTimeout || status >= 500
+}