feat: implement AI gateway phase one runtime
This commit is contained in:
@@ -0,0 +1,40 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func stringFromMap(values map[string]any, key string) string {
|
||||
value, _ := values[key].(string)
|
||||
return strings.TrimSpace(value)
|
||||
}
|
||||
|
||||
func boolFromMap(values map[string]any, key string) bool {
|
||||
value, _ := values[key].(bool)
|
||||
return value
|
||||
}
|
||||
|
||||
func intFromPolicy(values map[string]any, key string) int {
|
||||
switch value := values[key].(type) {
|
||||
case int:
|
||||
return value
|
||||
case float64:
|
||||
return int(math.Round(value))
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func floatFromAny(value any) float64 {
|
||||
switch typed := value.(type) {
|
||||
case int:
|
||||
return float64(typed)
|
||||
case int64:
|
||||
return float64(typed)
|
||||
case float64:
|
||||
return typed
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate, body map[string]any) []store.RateLimitReservation {
|
||||
out := make([]store.RateLimitReservation, 0)
|
||||
out = append(out, reservationsFromPolicy("platform_model", candidate.PlatformModelID, effectiveRateLimitPolicy(candidate), body)...)
|
||||
if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
|
||||
out = append(out, reservationsFromPolicy("user_group", group.ID, group.RateLimitPolicy, body)...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
|
||||
if hasRules(candidate.ModelRateLimitPolicy) {
|
||||
return candidate.ModelRateLimitPolicy
|
||||
}
|
||||
if hasRules(candidate.PlatformRateLimitPolicy) {
|
||||
return candidate.PlatformRateLimitPolicy
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func reservationsFromPolicy(scopeType string, scopeKey string, policy map[string]any, body map[string]any) []store.RateLimitReservation {
|
||||
if scopeKey == "" || !hasRules(policy) {
|
||||
return nil
|
||||
}
|
||||
rules, _ := policy["rules"].([]any)
|
||||
out := make([]store.RateLimitReservation, 0, len(rules))
|
||||
estimatedTokens := estimateRequestTokens(body)
|
||||
for _, rawRule := range rules {
|
||||
rule, _ := rawRule.(map[string]any)
|
||||
metric := strings.TrimSpace(stringFromMap(rule, "metric"))
|
||||
limit := floatFromAny(rule["limit"])
|
||||
amount := 1.0
|
||||
if strings.HasPrefix(metric, "tpm") {
|
||||
amount = float64(estimatedTokens)
|
||||
}
|
||||
out = append(out, store.RateLimitReservation{
|
||||
ScopeType: scopeType,
|
||||
ScopeKey: scopeKey,
|
||||
Metric: metric,
|
||||
Limit: limit,
|
||||
Amount: amount,
|
||||
WindowSeconds: int(floatFromAny(rule["windowSeconds"])),
|
||||
LeaseTTLSeconds: int(floatFromAny(rule["leaseTtlSeconds"])),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func hasRules(policy map[string]any) bool {
|
||||
rules, _ := policy["rules"].([]any)
|
||||
return len(rules) > 0
|
||||
}
|
||||
|
||||
func estimateRequestTokens(body map[string]any) int {
|
||||
text := ""
|
||||
if prompt := stringFromMap(body, "prompt"); prompt != "" {
|
||||
text += prompt
|
||||
}
|
||||
if input := stringFromMap(body, "input"); input != "" {
|
||||
text += input
|
||||
}
|
||||
if messages, ok := body["messages"].([]any); ok {
|
||||
for _, raw := range messages {
|
||||
message, _ := raw.(map[string]any)
|
||||
switch content := message["content"].(type) {
|
||||
case string:
|
||||
text += content
|
||||
case []any:
|
||||
for _, rawPart := range content {
|
||||
part, _ := rawPart.(map[string]any)
|
||||
text += stringFromMap(part, "text")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if text == "" {
|
||||
return 1
|
||||
}
|
||||
return len([]rune(text))/4 + 1
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
type EstimateResult struct {
|
||||
Items []any `json:"items"`
|
||||
Resolver string `json:"resolver"`
|
||||
}
|
||||
|
||||
func (s *Service) Estimate(ctx context.Context, kind string, model string, body map[string]any, user *auth.User) (EstimateResult, error) {
|
||||
candidates, err := s.store.ListModelCandidates(ctx, model, modelTypeFromKind(kind))
|
||||
if err != nil {
|
||||
return EstimateResult{}, err
|
||||
}
|
||||
candidate := candidates[0]
|
||||
usage := clients.Usage{InputTokens: estimateRequestTokens(body), OutputTokens: int(floatFromAny(body["max_tokens"]))}
|
||||
if usage.OutputTokens == 0 {
|
||||
usage.OutputTokens = 64
|
||||
}
|
||||
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
|
||||
response := clients.Response{Usage: usage, Result: map[string]any{"usage": map[string]any{
|
||||
"prompt_tokens": usage.InputTokens,
|
||||
"completion_tokens": usage.OutputTokens,
|
||||
"total_tokens": usage.TotalTokens,
|
||||
}}}
|
||||
return EstimateResult{
|
||||
Items: s.billings(ctx, user, kind, body, candidate, response, true),
|
||||
Resolver: "effective-pricing-v1",
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Service) billings(ctx context.Context, user *auth.User, kind string, body map[string]any, candidate store.RuntimeModelCandidate, response clients.Response, simulated bool) []any {
|
||||
config := effectiveBillingConfig(candidate)
|
||||
discount := effectiveDiscount(ctx, s.store, user, candidate)
|
||||
if modelTypeFromKind(kind) == "chat" {
|
||||
inputTokens := response.Usage.InputTokens
|
||||
outputTokens := response.Usage.OutputTokens
|
||||
if inputTokens == 0 && outputTokens == 0 {
|
||||
inputTokens = estimateRequestTokens(body)
|
||||
outputTokens = 1
|
||||
}
|
||||
inputAmount := roundPrice(float64(inputTokens) / 1000 * price(config, "textInputPer1k") * discount)
|
||||
outputAmount := roundPrice(float64(outputTokens) / 1000 * price(config, "textOutputPer1k") * discount)
|
||||
return []any{
|
||||
billingLine(candidate, "text_input", "1k_tokens", inputTokens, inputAmount, discount, simulated),
|
||||
billingLine(candidate, "text_output", "1k_tokens", outputTokens, outputAmount, discount, simulated),
|
||||
}
|
||||
}
|
||||
count := int(floatFromAny(body["n"]))
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
}
|
||||
resource := "image"
|
||||
baseKey := "imageBase"
|
||||
if kind == "images.edits" {
|
||||
resource = "image_edit"
|
||||
baseKey = "editBase"
|
||||
}
|
||||
amount := float64(count) * price(config, baseKey) * weighted(config, "qualityWeights", stringFromMap(body, "quality")) * weighted(config, "sizeWeights", stringFromMap(body, "size")) * discount
|
||||
return []any{billingLine(candidate, resource, "image", count, roundPrice(amount), discount, simulated)}
|
||||
}
|
||||
|
||||
func effectiveBillingConfig(candidate store.RuntimeModelCandidate) map[string]any {
|
||||
base := candidate.BaseBillingConfig
|
||||
if len(candidate.BillingConfig) > 0 {
|
||||
base = candidate.BillingConfig
|
||||
}
|
||||
if len(candidate.BillingConfigOverride) > 0 {
|
||||
base = mergeMap(base, candidate.BillingConfigOverride)
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func effectiveDiscount(ctx context.Context, db *store.Store, user *auth.User, candidate store.RuntimeModelCandidate) float64 {
|
||||
discount := candidate.DefaultDiscountFactor
|
||||
if candidate.DiscountFactor > 0 {
|
||||
discount = candidate.DiscountFactor
|
||||
}
|
||||
if discount <= 0 {
|
||||
discount = 1
|
||||
}
|
||||
if group, err := db.ResolveUserGroupPolicy(ctx, user); err == nil {
|
||||
groupDiscount := floatFromAny(group.BillingDiscountPolicy["discountFactor"])
|
||||
if groupDiscount > 0 {
|
||||
discount *= groupDiscount
|
||||
}
|
||||
}
|
||||
return discount
|
||||
}
|
||||
|
||||
func billingLine(candidate store.RuntimeModelCandidate, resourceType string, unit string, quantity any, amount float64, discount float64, simulated bool) map[string]any {
|
||||
return map[string]any{
|
||||
"model": candidate.ModelName,
|
||||
"modelAlias": candidate.ModelAlias,
|
||||
"provider": candidate.Provider,
|
||||
"platformId": candidate.PlatformID,
|
||||
"platformModelId": candidate.PlatformModelID,
|
||||
"resourceType": resourceType,
|
||||
"unit": unit,
|
||||
"quantity": quantity,
|
||||
"amount": amount,
|
||||
"currency": "resource",
|
||||
"discountFactor": discount,
|
||||
"simulated": simulated,
|
||||
}
|
||||
}
|
||||
|
||||
func price(config map[string]any, key string) float64 {
|
||||
value := floatFromAny(config[key])
|
||||
if value > 0 {
|
||||
return value
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func weighted(config map[string]any, key string, name string) float64 {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
return 1
|
||||
}
|
||||
weights, _ := config[key].(map[string]any)
|
||||
if value := floatFromAny(weights[name]); value > 0 {
|
||||
return value
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func roundPrice(value float64) float64 {
|
||||
return math.Round(value*1000000) / 1000000
|
||||
}
|
||||
|
||||
func mergeMap(base map[string]any, override map[string]any) map[string]any {
|
||||
out := map[string]any{}
|
||||
for key, value := range base {
|
||||
out[key] = value
|
||||
}
|
||||
for key, value := range override {
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
cfg config.Config
|
||||
store *store.Store
|
||||
logger *slog.Logger
|
||||
clients map[string]clients.Client
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Task store.GatewayTask
|
||||
Output map[string]any
|
||||
}
|
||||
|
||||
func New(cfg config.Config, db *store.Store, logger *slog.Logger) *Service {
|
||||
httpClient := &http.Client{Timeout: 120 * time.Second}
|
||||
return &Service{
|
||||
cfg: cfg,
|
||||
store: db,
|
||||
logger: logger,
|
||||
clients: map[string]clients.Client{
|
||||
"openai": clients.OpenAIClient{HTTPClient: httpClient},
|
||||
"gemini": clients.GeminiClient{HTTPClient: httpClient},
|
||||
"simulation": clients.SimulationClient{},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *auth.User) (Result, error) {
|
||||
modelType := modelTypeFromKind(task.Kind)
|
||||
body := normalizeRequest(task.Kind, task.Request)
|
||||
if err := validateRequest(task.Kind, body); err != nil {
|
||||
failed, finishErr := s.failTask(ctx, task.ID, "bad_request", err.Error(), task.RunMode == "simulation")
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, err
|
||||
}
|
||||
candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType)
|
||||
if err != nil {
|
||||
failed, finishErr := s.failTask(ctx, task.ID, "no_model_candidate", err.Error(), task.RunMode == "simulation")
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, err
|
||||
}
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, body); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := s.emit(ctx, task.ID, "task.progress", "running", "normalizing", 0.15, "request normalized", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
|
||||
maxAttempts := maxAttemptsForCandidates(candidates)
|
||||
var lastErr error
|
||||
for index, candidate := range candidates {
|
||||
if index >= maxAttempts {
|
||||
break
|
||||
}
|
||||
attemptNo := index + 1
|
||||
response, err := s.runCandidate(ctx, task, user, body, candidate, attemptNo)
|
||||
if err == nil {
|
||||
billings := s.billings(ctx, user, task.Kind, body, candidate, response, isSimulation(task, candidate))
|
||||
finished, finishErr := s.store.FinishTaskSuccess(ctx, task.ID, response.Result, billings)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
if err := s.emit(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "task completed", map[string]any{"result": response.Result, "billings": billings}, isSimulation(task, candidate)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return Result{Task: finished, Output: response.Result}, nil
|
||||
}
|
||||
lastErr = err
|
||||
retryable := clients.IsRetryable(err)
|
||||
if !retryable || !retryEnabled(candidate) || attemptNo >= maxAttempts {
|
||||
break
|
||||
}
|
||||
if err := s.emit(ctx, task.ID, "task.retrying", "running", "retry", 0.55, "retrying next client", map[string]any{"attempt": attemptNo, "error": err.Error()}, isSimulation(task, candidate)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
}
|
||||
code := clients.ErrorCode(lastErr)
|
||||
message := "task failed"
|
||||
if lastErr != nil {
|
||||
message = lastErr.Error()
|
||||
}
|
||||
failed, err := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation")
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, lastErr
|
||||
}
|
||||
|
||||
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, attemptNo int) (clients.Response, error) {
|
||||
simulated := isSimulation(task, candidate)
|
||||
if err := s.emit(ctx, task.ID, "task.attempt.started", "running", "submitting", 0.25, "client attempt started", map[string]any{"attempt": attemptNo, "clientId": candidate.ClientID}, simulated); err != nil {
|
||||
return clients.Response{}, err
|
||||
}
|
||||
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
|
||||
TaskID: task.ID,
|
||||
AttemptNo: attemptNo,
|
||||
PlatformID: candidate.PlatformID,
|
||||
PlatformModelID: candidate.PlatformModelID,
|
||||
ClientID: candidate.ClientID,
|
||||
QueueKey: candidate.QueueKey,
|
||||
Status: "running",
|
||||
Simulated: simulated,
|
||||
RequestSnapshot: body,
|
||||
})
|
||||
if err != nil {
|
||||
return clients.Response{}, err
|
||||
}
|
||||
reservations := s.rateLimitReservations(ctx, user, candidate, body)
|
||||
limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, reservations)
|
||||
if err != nil {
|
||||
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{AttemptID: attemptID, Status: "failed", Retryable: false, ErrorCode: "rate_limit", ErrorMessage: err.Error()})
|
||||
return clients.Response{}, &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: false}
|
||||
}
|
||||
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
|
||||
|
||||
if err := s.store.RecordClientAssignment(ctx, candidate); err != nil {
|
||||
return clients.Response{}, err
|
||||
}
|
||||
defer s.store.RecordClientRelease(context.WithoutCancel(ctx), candidate.ClientID, "")
|
||||
|
||||
client := s.clientFor(candidate, simulated)
|
||||
response, err := client.Run(ctx, clients.Request{
|
||||
Kind: task.Kind,
|
||||
ModelType: candidate.ModelType,
|
||||
Model: task.Model,
|
||||
Body: body,
|
||||
Candidate: candidate,
|
||||
Stream: boolFromMap(body, "stream"),
|
||||
})
|
||||
if err != nil {
|
||||
retryable := clients.IsRetryable(err)
|
||||
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
||||
AttemptID: attemptID,
|
||||
Status: "failed",
|
||||
Retryable: retryable,
|
||||
ErrorCode: clients.ErrorCode(err),
|
||||
ErrorMessage: err.Error(),
|
||||
})
|
||||
_ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable}, simulated)
|
||||
return clients.Response{}, err
|
||||
}
|
||||
uploadedResult, err := s.uploadGeneratedAssets(ctx, response.Result)
|
||||
if err != nil {
|
||||
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
||||
AttemptID: attemptID,
|
||||
Status: "failed",
|
||||
Retryable: clients.IsRetryable(err),
|
||||
ErrorCode: clients.ErrorCode(err),
|
||||
ErrorMessage: err.Error(),
|
||||
})
|
||||
return clients.Response{}, err
|
||||
}
|
||||
response.Result = uploadedResult
|
||||
for _, progress := range response.Progress {
|
||||
if err := s.emit(ctx, task.ID, "task.progress", "running", progress.Phase, progress.Progress, progress.Message, progress.Payload, simulated); err != nil {
|
||||
return clients.Response{}, err
|
||||
}
|
||||
}
|
||||
if err := s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
||||
AttemptID: attemptID,
|
||||
Status: "succeeded",
|
||||
ResponseSnapshot: response.Result,
|
||||
}); err != nil {
|
||||
return clients.Response{}, err
|
||||
}
|
||||
return response, nil
|
||||
}
|
||||
|
||||
func (s *Service) clientFor(candidate store.RuntimeModelCandidate, simulated bool) clients.Client {
|
||||
if simulated {
|
||||
return s.clients["simulation"]
|
||||
}
|
||||
key := strings.ToLower(candidate.Provider)
|
||||
if client, ok := s.clients[key]; ok {
|
||||
return client
|
||||
}
|
||||
return s.clients["openai"]
|
||||
}
|
||||
|
||||
func (s *Service) failTask(ctx context.Context, taskID string, code string, message string, simulated bool) (store.GatewayTask, error) {
|
||||
failed, err := s.store.FinishTaskFailure(ctx, taskID, code, message)
|
||||
if err != nil {
|
||||
return store.GatewayTask{}, err
|
||||
}
|
||||
if eventErr := s.emit(ctx, taskID, "task.failed", "failed", "failed", 1, message, map[string]any{"code": code}, simulated); eventErr != nil {
|
||||
return store.GatewayTask{}, eventErr
|
||||
}
|
||||
return failed, nil
|
||||
}
|
||||
|
||||
func (s *Service) emit(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) error {
|
||||
event, err := s.store.AddTaskEvent(ctx, taskID, eventType, status, phase, progress, message, payload, simulated)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if s.cfg.TaskProgressCallbackEnabled {
|
||||
return s.store.QueueTaskCallback(ctx, event, s.cfg.TaskProgressCallbackURL)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func modelTypeFromKind(kind string) string {
|
||||
switch kind {
|
||||
case "chat.completions", "responses":
|
||||
return "chat"
|
||||
case "images.generations", "images.edits":
|
||||
return "image"
|
||||
default:
|
||||
return "task"
|
||||
}
|
||||
}
|
||||
|
||||
func isSimulation(task store.GatewayTask, candidate store.RuntimeModelCandidate) bool {
|
||||
if task.RunMode == "simulation" {
|
||||
return true
|
||||
}
|
||||
return stringFromMap(candidate.Credentials, "mode") == "simulation" || boolFromMap(candidate.PlatformConfig, "testMode")
|
||||
}
|
||||
|
||||
func retryEnabled(candidate store.RuntimeModelCandidate) bool {
|
||||
if enabled, ok := candidate.ModelRetryPolicy["enabled"].(bool); ok {
|
||||
return enabled
|
||||
}
|
||||
if enabled, ok := candidate.PlatformRetryPolicy["enabled"].(bool); ok {
|
||||
return enabled
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func maxAttemptsForCandidates(candidates []store.RuntimeModelCandidate) int {
|
||||
if len(candidates) == 0 {
|
||||
return 0
|
||||
}
|
||||
maxAttempts := len(candidates)
|
||||
for _, candidate := range candidates {
|
||||
if value := intFromPolicy(candidate.ModelRetryPolicy, "maxAttempts"); value > 0 && value < maxAttempts {
|
||||
maxAttempts = value
|
||||
}
|
||||
if value := intFromPolicy(candidate.PlatformRetryPolicy, "maxAttempts"); value > 0 && value < maxAttempts {
|
||||
maxAttempts = value
|
||||
}
|
||||
}
|
||||
if maxAttempts <= 0 {
|
||||
return 1
|
||||
}
|
||||
return maxAttempts
|
||||
}
|
||||
|
||||
func normalizeRequest(kind string, body map[string]any) map[string]any {
|
||||
out := map[string]any{}
|
||||
for key, value := range body {
|
||||
out[key] = value
|
||||
}
|
||||
if kind == "responses" && out["messages"] == nil && out["input"] != nil {
|
||||
out["messages"] = []any{map[string]any{"role": "user", "content": out["input"]}}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func validateRequest(kind string, body map[string]any) error {
|
||||
switch kind {
|
||||
case "chat.completions":
|
||||
if body["messages"] == nil {
|
||||
return errors.New("messages is required")
|
||||
}
|
||||
case "responses":
|
||||
if body["input"] == nil && body["messages"] == nil {
|
||||
return errors.New("input or messages is required")
|
||||
}
|
||||
case "images.generations", "images.edits":
|
||||
if strings.TrimSpace(stringFromMap(body, "prompt")) == "" {
|
||||
return errors.New("prompt is required")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
||||
)
|
||||
|
||||
func (s *Service) uploadGeneratedAssets(ctx context.Context, result map[string]any) (map[string]any, error) {
|
||||
if s.cfg.ServerMainBaseURL == "" || s.cfg.ServerMainInternalToken == "" {
|
||||
return result, nil
|
||||
}
|
||||
data, _ := result["data"].([]any)
|
||||
if len(data) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
next := map[string]any{}
|
||||
for key, value := range result {
|
||||
next[key] = value
|
||||
}
|
||||
nextData := make([]any, 0, len(data))
|
||||
for index, rawItem := range data {
|
||||
item, _ := rawItem.(map[string]any)
|
||||
if item == nil {
|
||||
nextData = append(nextData, rawItem)
|
||||
continue
|
||||
}
|
||||
b64 := stringFromMap(item, "b64_json")
|
||||
if b64 == "" {
|
||||
nextData = append(nextData, rawItem)
|
||||
continue
|
||||
}
|
||||
upload, err := s.uploadBase64Image(ctx, b64, index)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
merged := map[string]any{}
|
||||
for key, value := range item {
|
||||
if key != "b64_json" {
|
||||
merged[key] = value
|
||||
}
|
||||
}
|
||||
merged["upload"] = upload
|
||||
if urlValue, ok := upload["url"].(string); ok && urlValue != "" {
|
||||
merged["url"] = urlValue
|
||||
}
|
||||
nextData = append(nextData, merged)
|
||||
}
|
||||
next["data"] = nextData
|
||||
return next, nil
|
||||
}
|
||||
|
||||
func (s *Service) uploadBase64Image(ctx context.Context, b64 string, index int) (map[string]any, error) {
|
||||
payload, err := base64.StdEncoding.DecodeString(stripDataURLPrefix(b64))
|
||||
if err != nil {
|
||||
return nil, &clients.ClientError{Code: "upload_decode_failed", Message: err.Error(), Retryable: false}
|
||||
}
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
fileWriter, err := writer.CreateFormFile("file", fmt.Sprintf("gateway-result-%d.png", index+1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if _, err := fileWriter.Write(payload); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_ = writer.WriteField("source", "ai-gateway")
|
||||
if err := writer.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(s.cfg.ServerMainBaseURL, "/")+"/v1/files/upload", &body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
req.Header.Set("Authorization", "Bearer "+s.cfg.ServerMainInternalToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, &clients.ClientError{Code: "upload_network", Message: err.Error(), Retryable: true}
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var decoded map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&decoded); err != nil {
|
||||
return nil, &clients.ClientError{Code: "upload_invalid_response", Message: err.Error(), Retryable: false}
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, &clients.ClientError{Code: "upload_failed", Message: "server-main upload failed", StatusCode: resp.StatusCode, Retryable: resp.StatusCode >= 500}
|
||||
}
|
||||
return decoded, nil
|
||||
}
|
||||
|
||||
func stripDataURLPrefix(value string) string {
|
||||
if index := strings.Index(value, ","); strings.HasPrefix(value, "data:") && index >= 0 {
|
||||
return value[index+1:]
|
||||
}
|
||||
return value
|
||||
}
|
||||
Reference in New Issue
Block a user