BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement AI gateway phase one runtime

Browse Source
This commit is contained in:
wangbo
2026-05-09 21:18:32 +08:00
parent a5e66e79cd
commit fdcdcd477b
46 changed files with 5678 additions and 768 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/clients/clients_test.go
+120
View File
@@ -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
}
apps/api/internal/clients/gemini.go
+170
View File
@@ -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}
}
apps/api/internal/clients/helpers.go
+133
View File
@@ -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()
}
apps/api/internal/clients/openai.go
+87
View File
@@ -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}},
}
}
apps/api/internal/clients/simulation.go
+119
View File
@@ -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
}
apps/api/internal/clients/types.go
+72
View File
@@ -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
}
apps/api/internal/httpapi/catalog_handlers.go
+161
View File
@@ -0,0 +1,161 @@
package httpapi
import (
"encoding/json"
"net/http"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func (s *Server) listCatalogProviders(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListCatalogProviders(r.Context())
if err != nil {
s.logger.Error("list catalog providers failed", "error", err)
writeError(w, http.StatusInternalServerError, "list catalog providers failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) createCatalogProvider(w http.ResponseWriter, r *http.Request) {
var input store.CatalogProviderInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if strings.TrimSpace(input.ProviderKey) == "" || strings.TrimSpace(input.DisplayName) == "" {
writeError(w, http.StatusBadRequest, "providerKey and displayName are required")
return
}
item, err := s.store.CreateCatalogProvider(r.Context(), input)
if err != nil {
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "provider key or code already exists")
return
}
s.logger.Error("create catalog provider failed", "error", err)
writeError(w, http.StatusInternalServerError, "create catalog provider failed")
return
}
writeJSON(w, http.StatusCreated, item)
}
func (s *Server) updateCatalogProvider(w http.ResponseWriter, r *http.Request) {
var input store.CatalogProviderInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if strings.TrimSpace(input.ProviderKey) == "" || strings.TrimSpace(input.DisplayName) == "" {
writeError(w, http.StatusBadRequest, "providerKey and displayName are required")
return
}
item, err := s.store.UpdateCatalogProvider(r.Context(), r.PathValue("providerID"), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "catalog provider not found")
return
}
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "provider key or code already exists")
return
}
s.logger.Error("update catalog provider failed", "error", err)
writeError(w, http.StatusInternalServerError, "update catalog provider failed")
return
}
writeJSON(w, http.StatusOK, item)
}
func (s *Server) deleteCatalogProvider(w http.ResponseWriter, r *http.Request) {
if err := s.store.DeleteCatalogProvider(r.Context(), r.PathValue("providerID")); err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "catalog provider not found")
return
}
s.logger.Error("delete catalog provider failed", "error", err)
writeError(w, http.StatusInternalServerError, "delete catalog provider failed")
return
}
w.WriteHeader(http.StatusNoContent)
}
func (s *Server) listBaseModels(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListBaseModels(r.Context())
if err != nil {
s.logger.Error("list base models failed", "error", err)
writeError(w, http.StatusInternalServerError, "list base models failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) createBaseModel(w http.ResponseWriter, r *http.Request) {
var input store.BaseModelInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if !validBaseModelInput(input) {
writeError(w, http.StatusBadRequest, "providerKey, providerModelName and modelType are required")
return
}
item, err := s.store.CreateBaseModel(r.Context(), input)
if err != nil {
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "canonical model key already exists")
return
}
s.logger.Error("create base model failed", "error", err)
writeError(w, http.StatusInternalServerError, "create base model failed")
return
}
writeJSON(w, http.StatusCreated, item)
}
func (s *Server) updateBaseModel(w http.ResponseWriter, r *http.Request) {
var input store.BaseModelInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if !validBaseModelInput(input) {
writeError(w, http.StatusBadRequest, "providerKey, providerModelName and modelType are required")
return
}
item, err := s.store.UpdateBaseModel(r.Context(), r.PathValue("baseModelID"), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "base model not found")
return
}
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "canonical model key already exists")
return
}
s.logger.Error("update base model failed", "error", err)
writeError(w, http.StatusInternalServerError, "update base model failed")
return
}
writeJSON(w, http.StatusOK, item)
}
func (s *Server) deleteBaseModel(w http.ResponseWriter, r *http.Request) {
if err := s.store.DeleteBaseModel(r.Context(), r.PathValue("baseModelID")); err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "base model not found")
return
}
s.logger.Error("delete base model failed", "error", err)
writeError(w, http.StatusInternalServerError, "delete base model failed")
return
}
w.WriteHeader(http.StatusNoContent)
}
func validBaseModelInput(input store.BaseModelInput) bool {
return strings.TrimSpace(input.ProviderKey) != "" &&
strings.TrimSpace(input.ProviderModelName) != "" &&
strings.TrimSpace(input.ModelType) != ""
}
apps/api/internal/httpapi/core_flow_integration_test.go
+182 -7
View File
@@ -37,12 +37,14 @@ func TestCoreLocalFlow(t *testing.T) {
defer db.Close()
handler := NewServer(config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
TaskProgressCallbackEnabled: true,
TaskProgressCallbackURL: "http://callback.local/task-progress",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
server := httptest.NewServer(handler)
defer server.Close()
@@ -117,6 +119,9 @@ func TestCoreLocalFlow(t *testing.T) {
t.Fatalf("connect test pool: %v", err)
}
defer testPool.Close()
if _, err := testPool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote smoke user: %v", err)
}
inviteCode := "INVITE-" + suffixText
if _, err := testPool.Exec(ctx, `
@@ -167,7 +172,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
PlatformKey string `json:"platformKey"`
Status string `json:"status"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/platforms", loginResponse.AccessToken, map[string]any{
doJSON(t, server.URL, http.MethodPost, "/api/v1/platforms", apiKeyResponse.Secret, map[string]any{
"provider": "openai",
"platformKey": "openai-smoke-" + suffixText,
"name": "OpenAI Smoke",
@@ -180,6 +185,47 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatalf("unexpected platform response: %+v", platform)
}
var baseModels struct {
Items []struct {
ID string `json:"id"`
CanonicalModelKey string `json:"canonicalModelKey"`
ProviderModelName string `json:"providerModelName"`
ModelType string `json:"modelType"`
} `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/v1/catalog/base-models", apiKeyResponse.Secret, nil, http.StatusOK, &baseModels)
if len(baseModels.Items) < 300 {
t.Fatalf("server-main seed should include the migrated base model catalog: got %d", len(baseModels.Items))
}
baseModelInput := map[string]any{
"providerKey": "openai",
"canonicalModelKey": "openai:smoke-base-" + suffixText,
"providerModelName": "smoke-base-" + suffixText,
"modelType": "text_generate",
"displayName": "Smoke Base Model",
"capabilities": map[string]any{"originalTypes": []string{"text_generate"}},
"metadata": map[string]any{"source": "test"},
}
var createdBaseModel struct {
ID string `json:"id"`
CanonicalModelKey string `json:"canonicalModelKey"`
DisplayName string `json:"displayName"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/catalog/base-models", apiKeyResponse.Secret, baseModelInput, http.StatusCreated, &createdBaseModel)
if createdBaseModel.ID == "" || createdBaseModel.CanonicalModelKey != baseModelInput["canonicalModelKey"] {
t.Fatalf("unexpected created base model: %+v", createdBaseModel)
}
baseModelInput["displayName"] = "Smoke Base Model Updated"
var updatedBaseModel struct {
DisplayName string `json:"displayName"`
}
doJSON(t, server.URL, http.MethodPatch, "/api/v1/catalog/base-models/"+createdBaseModel.ID, apiKeyResponse.Secret, baseModelInput, http.StatusOK, &updatedBaseModel)
if updatedBaseModel.DisplayName != "Smoke Base Model Updated" {
t.Fatalf("unexpected updated base model: %+v", updatedBaseModel)
}
doJSON(t, server.URL, http.MethodDelete, "/api/v1/catalog/base-models/"+createdBaseModel.ID, apiKeyResponse.Secret, nil, http.StatusNoContent, nil)
var taskResponse struct {
Task struct {
ID string `json:"id"`
@@ -198,6 +244,109 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
t.Fatalf("unexpected task response: %+v", taskResponse.Task)
}
var compatChat map[string]any
doJSON(t, server.URL, http.MethodPost, "/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": "gpt-4o-mini",
"runMode": "simulation",
"messages": []map[string]any{{"role": "user", "content": "ping"}},
"simulation": true,
}, http.StatusOK, &compatChat)
if compatChat["object"] != "chat.completion" {
t.Fatalf("unexpected compatible chat response: %+v", compatChat)
}
var imageResponse struct {
Task struct {
ID string `json:"id"`
Status string `json:"status"`
Result map[string]any `json:"result"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
"model": "gpt-image-1",
"runMode": "simulation",
"prompt": "a tiny gateway console",
"size": "1024x1024",
"quality": "medium",
"simulation": true,
}, http.StatusAccepted, &imageResponse)
if imageResponse.Task.Status != "succeeded" || imageResponse.Task.Result["id"] == "" {
t.Fatalf("unexpected image generation task: %+v", imageResponse.Task)
}
var imageEditResponse struct {
Task struct {
ID string `json:"id"`
Status string `json:"status"`
Result map[string]any `json:"result"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/edits", apiKeyResponse.Secret, map[string]any{
"model": "gpt-image-1",
"runMode": "simulation",
"prompt": "replace background with clean studio light",
"image": "https://example.com/source.png",
"mask": "https://example.com/mask.png",
"simulation": true,
}, http.StatusAccepted, &imageEditResponse)
if imageEditResponse.Task.Status != "succeeded" || imageEditResponse.Task.Result["id"] == "" {
t.Fatalf("unexpected image edit task: %+v", imageEditResponse.Task)
}
failoverModel := "phase1-failover-" + suffixText
var failedPlatform struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/platforms", apiKeyResponse.Secret, map[string]any{
"provider": "openai",
"platformKey": "openai-fail-" + suffixText,
"name": "OpenAI Retryable Failure",
"baseUrl": "https://api.openai.com/v1",
"authType": "bearer",
"credentials": map[string]any{"mode": "simulation", "simulationFailure": "retryable_failure"},
"priority": 10,
}, http.StatusCreated, &failedPlatform)
var successPlatform struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/platforms", apiKeyResponse.Secret, map[string]any{
"provider": "openai",
"platformKey": "openai-success-" + suffixText,
"name": "OpenAI Retry Success",
"baseUrl": "https://api.openai.com/v1",
"authType": "bearer",
"credentials": map[string]any{"mode": "simulation"},
"priority": 20,
}, http.StatusCreated, &successPlatform)
for _, platformID := range []string{failedPlatform.ID, successPlatform.ID} {
var platformModel map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/platforms/"+platformID+"/models", apiKeyResponse.Secret, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"modelName": failoverModel,
"modelAlias": failoverModel,
"modelType": "chat",
"displayName": "Failover Smoke",
"retryPolicy": map[string]any{"enabled": true, "maxAttempts": 2},
}, http.StatusCreated, &platformModel)
if platformModel["id"] == "" {
t.Fatalf("platform model was not created: %+v", platformModel)
}
}
var failoverTask struct {
Task struct {
ID string `json:"id"`
Status string `json:"status"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, map[string]any{
"model": failoverModel,
"runMode": "simulation",
"messages": []map[string]any{{"role": "user", "content": "retry please"}},
}, http.StatusAccepted, &failoverTask)
if failoverTask.Task.Status != "succeeded" {
t.Fatalf("failover task should succeed through second client: %+v", failoverTask.Task)
}
var taskDetail struct {
ID string `json:"id"`
Status string `json:"status"`
@@ -222,6 +371,32 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
if resp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("task.completed")) {
t.Fatalf("unexpected events response status=%d body=%s", resp.StatusCode, string(body))
}
if !bytes.Contains(body, []byte("task.progress")) {
t.Fatalf("events response should include progress events body=%s", string(body))
}
req, err = http.NewRequest(http.MethodGet, server.URL+"/api/v1/tasks/"+failoverTask.Task.ID+"/events", nil)
if err != nil {
t.Fatalf("build failover events request: %v", err)
}
req.Header.Set("Authorization", "Bearer "+apiKeyResponse.Secret)
resp, err = http.DefaultClient.Do(req)
if err != nil {
t.Fatalf("failover events request: %v", err)
}
defer resp.Body.Close()
body, _ = io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK || !bytes.Contains(body, []byte("task.retrying")) {
t.Fatalf("failover events should include retrying event status=%d body=%s", resp.StatusCode, string(body))
}
var callbackRows int
if err := testPool.QueryRow(ctx, `SELECT count(*) FROM gateway_task_callback_outbox WHERE task_id = $1::uuid`, taskResponse.Task.ID).Scan(&callbackRows); err != nil {
t.Fatalf("read callback outbox: %v", err)
}
if callbackRows == 0 {
t.Fatal("task progress callback outbox should receive events")
}
}
func TestOriginAllowedSupportsCommaSeparatedOrigins(t *testing.T) {
apps/api/internal/httpapi/handlers.go
+81 -27
View File
@@ -164,6 +164,32 @@ func (s *Server) createPlatform(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusCreated, platform)
}
func (s *Server) createPlatformModel(w http.ResponseWriter, r *http.Request) {
var input store.CreatePlatformModelInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if pathPlatformID := r.PathValue("platformID"); pathPlatformID != "" {
input.PlatformID = pathPlatformID
}
if input.PlatformID == "" {
writeError(w, http.StatusBadRequest, "platformId is required")
return
}
model, err := s.store.CreatePlatformModel(r.Context(), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "base model not found")
return
}
s.logger.Error("create platform model failed", "error", err)
writeError(w, http.StatusInternalServerError, "create platform model failed")
return
}
writeJSON(w, http.StatusCreated, model)
}
func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
models, err := s.store.ListModels(r.Context())
if err != nil {
@@ -174,26 +200,6 @@ func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"items": models})
}
func (s *Server) listCatalogProviders(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListCatalogProviders(r.Context())
if err != nil {
s.logger.Error("list catalog providers failed", "error", err)
writeError(w, http.StatusInternalServerError, "list catalog providers failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) listBaseModels(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListBaseModels(r.Context())
if err != nil {
s.logger.Error("list base models failed", "error", err)
writeError(w, http.StatusInternalServerError, "list base models failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) listPricingRules(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListPricingRules(r.Context())
if err != nil {
@@ -285,16 +291,32 @@ func (s *Server) disableAPIKey(w http.ResponseWriter, r *http.Request) {
}
func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context())
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
writeJSON(w, http.StatusOK, map[string]any{
"items": []any{},
"resolver": "effective-pricing-placeholder",
"request": body,
})
model, _ := body["model"].(string)
kind, _ := body["kind"].(string)
if kind == "" {
kind = "chat.completions"
}
if model == "" {
writeError(w, http.StatusBadRequest, "model is required")
return
}
estimate, err := s.runner.Estimate(r.Context(), kind, model, body, user)
if err != nil {
if errors.Is(err, store.ErrNoModelCandidate) {
writeError(w, http.StatusNotFound, "no enabled platform model matches request")
return
}
s.logger.Error("estimate pricing failed", "error", err)
writeError(w, http.StatusInternalServerError, "estimate pricing failed")
return
}
writeJSON(w, http.StatusOK, estimate)
}
func (s *Server) listRateLimitWindows(w http.ResponseWriter, r *http.Request) {
@@ -307,7 +329,7 @@ func (s *Server) listRateLimitWindows(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) createTask(kind string) http.Handler {
func (s *Server) createTask(kind string, compatible bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
@@ -337,9 +359,25 @@ func (s *Server) createTask(kind string) http.Handler {
writeError(w, http.StatusInternalServerError, "create task failed")
return
}
result, runErr := s.runner.Execute(r.Context(), task, user)
if compatible {
if runErr != nil {
writeError(w, statusFromRunError(runErr), runErr.Error())
return
}
if boolValue(body, "stream") {
writeCompatibleStream(w, kind, model, result.Output)
return
}
writeJSON(w, http.StatusOK, result.Output)
return
}
if runErr != nil {
s.logger.Warn("task completed with failure", "kind", kind, "taskId", task.ID, "error", runErr)
}
writeJSON(w, http.StatusAccepted, map[string]any{
"task": task,
"task": result.Task,
"next": map[string]string{
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
@@ -348,6 +386,22 @@ func (s *Server) createTask(kind string) http.Handler {
})
}
func statusFromRunError(err error) int {
switch {
case errors.Is(err, store.ErrNoModelCandidate):
return http.StatusNotFound
case errors.Is(err, store.ErrRateLimited):
return http.StatusTooManyRequests
default:
return http.StatusBadGateway
}
}
func boolValue(body map[string]any, key string) bool {
value, _ := body[key].(bool)
return value
}
func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
if err == nil {
apps/api/internal/httpapi/pricing_handlers.go
+94
View File
@@ -0,0 +1,94 @@
package httpapi
import (
"encoding/json"
"net/http"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func (s *Server) listPricingRuleSets(w http.ResponseWriter, r *http.Request) {
items, err := s.store.ListPricingRuleSets(r.Context())
if err != nil {
s.logger.Error("list pricing rule sets failed", "error", err)
writeError(w, http.StatusInternalServerError, "list pricing rule sets failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
func (s *Server) createPricingRuleSet(w http.ResponseWriter, r *http.Request) {
var input store.PricingRuleSetInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if !validPricingRuleSetInput(input) {
writeError(w, http.StatusBadRequest, "ruleSetKey, name and at least one rule are required")
return
}
item, err := s.store.CreatePricingRuleSet(r.Context(), input)
if err != nil {
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "pricing rule set key already exists")
return
}
s.logger.Error("create pricing rule set failed", "error", err)
writeError(w, http.StatusInternalServerError, "create pricing rule set failed")
return
}
writeJSON(w, http.StatusCreated, item)
}
func (s *Server) updatePricingRuleSet(w http.ResponseWriter, r *http.Request) {
var input store.PricingRuleSetInput
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if !validPricingRuleSetInput(input) {
writeError(w, http.StatusBadRequest, "ruleSetKey, name and at least one rule are required")
return
}
item, err := s.store.UpdatePricingRuleSet(r.Context(), r.PathValue("ruleSetID"), input)
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "pricing rule set not found")
return
}
if store.IsUniqueViolation(err) {
writeError(w, http.StatusConflict, "pricing rule set key already exists")
return
}
s.logger.Error("update pricing rule set failed", "error", err)
writeError(w, http.StatusInternalServerError, "update pricing rule set failed")
return
}
writeJSON(w, http.StatusOK, item)
}
func (s *Server) deletePricingRuleSet(w http.ResponseWriter, r *http.Request) {
if err := s.store.DeletePricingRuleSet(r.Context(), r.PathValue("ruleSetID")); err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "pricing rule set not found")
return
}
s.logger.Error("delete pricing rule set failed", "error", err)
writeError(w, http.StatusInternalServerError, "delete pricing rule set failed")
return
}
w.WriteHeader(http.StatusNoContent)
}
func validPricingRuleSetInput(input store.PricingRuleSetInput) bool {
if strings.TrimSpace(input.RuleSetKey) == "" || strings.TrimSpace(input.Name) == "" || len(input.Rules) == 0 {
return false
}
for _, rule := range input.Rules {
if strings.TrimSpace(rule.ResourceType) == "" || strings.TrimSpace(rule.Unit) == "" {
return false
}
}
return true
}
apps/api/internal/httpapi/server.go
+30 -3
View File
@@ -7,6 +7,7 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
@@ -14,6 +15,7 @@ type Server struct {
cfg config.Config
store *store.Store
auth *auth.Authenticator
runner *runner.Service
logger *slog.Logger
}
@@ -22,6 +24,7 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
cfg: cfg,
store: db,
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
runner: runner.New(cfg, db, logger),
logger: logger,
}
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
@@ -33,8 +36,16 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
mux.Handle("POST /api/v1/auth/register", server.auth.Require(auth.PermissionPublic, http.HandlerFunc(server.register)))
mux.Handle("POST /api/v1/auth/login", server.auth.Require(auth.PermissionPublic, http.HandlerFunc(server.login)))
mux.Handle("GET /api/v1/me", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.me)))
mux.Handle("GET /api/v1/public/catalog/providers", server.auth.Require(auth.PermissionPublic, http.HandlerFunc(server.listCatalogProviders)))
mux.Handle("GET /api/v1/public/catalog/base-models", server.auth.Require(auth.PermissionPublic, http.HandlerFunc(server.listBaseModels)))
mux.Handle("GET /api/v1/catalog/providers", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listCatalogProviders)))
mux.Handle("POST /api/v1/catalog/providers", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.createCatalogProvider)))
mux.Handle("PATCH /api/v1/catalog/providers/{providerID}", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.updateCatalogProvider)))
mux.Handle("DELETE /api/v1/catalog/providers/{providerID}", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.deleteCatalogProvider)))
mux.Handle("GET /api/v1/catalog/base-models", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listBaseModels)))
mux.Handle("POST /api/v1/catalog/base-models", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.createBaseModel)))
mux.Handle("PATCH /api/v1/catalog/base-models/{baseModelID}", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.updateBaseModel)))
mux.Handle("DELETE /api/v1/catalog/base-models/{baseModelID}", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.deleteBaseModel)))
mux.Handle("GET /api/v1/tenants", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listTenants)))
mux.Handle("GET /api/v1/users", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listUsers)))
mux.Handle("GET /api/v1/user-groups", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listUserGroups)))
@@ -42,16 +53,32 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
mux.Handle("POST /api/v1/api-keys", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.createAPIKey)))
mux.Handle("PATCH /api/v1/api-keys/{apiKeyID}/disable", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.disableAPIKey)))
mux.Handle("GET /api/v1/pricing/rules", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPricingRules)))
mux.Handle("GET /api/v1/pricing/rule-sets", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPricingRuleSets)))
mux.Handle("POST /api/v1/pricing/rule-sets", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.createPricingRuleSet)))
mux.Handle("PATCH /api/v1/pricing/rule-sets/{ruleSetID}", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.updatePricingRuleSet)))
mux.Handle("DELETE /api/v1/pricing/rule-sets/{ruleSetID}", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.deletePricingRuleSet)))
mux.Handle("POST /api/v1/pricing/estimate", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.estimatePricing)))
mux.Handle("GET /api/v1/platforms", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listPlatforms)))
mux.Handle("POST /api/v1/platforms", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.createPlatform)))
mux.Handle("POST /api/v1/platforms/{platformID}/models", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.createPlatformModel)))
mux.Handle("POST /api/v1/platform-models", server.auth.Require(auth.PermissionManager, http.HandlerFunc(server.createPlatformModel)))
mux.Handle("GET /api/v1/models", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listModels)))
mux.Handle("GET /api/v1/runtime/rate-limit-windows", server.auth.Require(auth.PermissionPower, http.HandlerFunc(server.listRateLimitWindows)))
mux.Handle("POST /api/v1/chat/completions", server.auth.Require(auth.PermissionBasic, server.createTask("chat.completions")))
mux.Handle("POST /api/v1/images/generations", server.auth.Require(auth.PermissionBasic, server.createTask("images.generations")))
mux.Handle("POST /api/v1/videos/generations", server.auth.Require(auth.PermissionBasic, server.createTask("videos.generations")))
mux.Handle("POST /api/v1/chat/completions", server.auth.Require(auth.PermissionBasic, server.createTask("chat.completions", false)))
mux.Handle("POST /api/v1/responses", server.auth.Require(auth.PermissionBasic, server.createTask("responses", false)))
mux.Handle("POST /api/v1/images/generations", server.auth.Require(auth.PermissionBasic, server.createTask("images.generations", false)))
mux.Handle("POST /api/v1/images/edits", server.auth.Require(auth.PermissionBasic, server.createTask("images.edits", false)))
mux.Handle("POST /api/v1/videos/generations", server.auth.Require(auth.PermissionBasic, server.createTask("videos.generations", false)))
mux.Handle("GET /api/v1/tasks/{taskID}", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.getTask)))
mux.Handle("GET /api/v1/tasks/{taskID}/events", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.taskEvents)))
mux.Handle("POST /chat/completions", server.auth.Require(auth.PermissionBasic, server.createTask("chat.completions", true)))
mux.Handle("POST /v1/chat/completions", server.auth.Require(auth.PermissionBasic, server.createTask("chat.completions", true)))
mux.Handle("POST /responses", server.auth.Require(auth.PermissionBasic, server.createTask("responses", true)))
mux.Handle("POST /v1/responses", server.auth.Require(auth.PermissionBasic, server.createTask("responses", true)))
mux.Handle("POST /images/generations", server.auth.Require(auth.PermissionBasic, server.createTask("images.generations", true)))
mux.Handle("POST /v1/images/generations", server.auth.Require(auth.PermissionBasic, server.createTask("images.generations", true)))
mux.Handle("POST /images/edits", server.auth.Require(auth.PermissionBasic, server.createTask("images.edits", true)))
mux.Handle("POST /v1/images/edits", server.auth.Require(auth.PermissionBasic, server.createTask("images.edits", true)))
return server.recover(server.cors(mux))
}
apps/api/internal/httpapi/streaming.go
+45
View File
@@ -0,0 +1,45 @@
package httpapi
import "net/http"
func writeCompatibleStream(w http.ResponseWriter, kind string, model string, output map[string]any) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
content := extractOutputText(output)
if content == "" {
content = "done"
}
if kind == "responses" {
sendSSE(w, "response.output_text.delta", map[string]any{"type": "response.output_text.delta", "delta": content})
sendSSE(w, "response.completed", map[string]any{"type": "response.completed", "response": output})
return
}
sendSSE(w, "message", map[string]any{
"id": output["id"],
"object": "chat.completion.chunk",
"model": model,
"choices": []any{map[string]any{"index": 0, "delta": map[string]any{"content": content}, "finish_reason": nil}},
})
sendSSE(w, "message", map[string]any{
"id": output["id"],
"object": "chat.completion.chunk",
"model": model,
"choices": []any{map[string]any{"index": 0, "delta": map[string]any{}, "finish_reason": "stop"}},
})
}
func extractOutputText(output map[string]any) string {
if text, ok := output["output_text"].(string); ok {
return text
}
choices, _ := output["choices"].([]any)
for _, rawChoice := range choices {
choice, _ := rawChoice.(map[string]any)
message, _ := choice["message"].(map[string]any)
if content, ok := message["content"].(string); ok {
return content
}
}
return ""
}
apps/api/internal/runner/helpers.go
+40
View File
@@ -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
}
}
apps/api/internal/runner/limits.go
+89
View File
@@ -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
}
apps/api/internal/runner/pricing.go
+148
View File
@@ -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
}
apps/api/internal/runner/service.go
+295
View File
@@ -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
}
apps/api/internal/runner/upload.go
+104
View File
@@ -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
}
apps/api/internal/store/base_models.go
+190
View File
@@ -0,0 +1,190 @@
package store
import (
"context"
"encoding/json"
"strings"
"github.com/jackc/pgx/v5"
)
const baseModelColumns = `
id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy, metadata, pricing_version,
status, created_at, updated_at`
type BaseModelInput struct {
ProviderKey string `json:"providerKey"`
CanonicalModelKey string `json:"canonicalModelKey"`
ProviderModelName string `json:"providerModelName"`
ModelType string `json:"modelType"`
DisplayName string `json:"displayName"`
Capabilities map[string]any `json:"capabilities"`
BaseBillingConfig map[string]any `json:"baseBillingConfig"`
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy"`
Metadata map[string]any `json:"metadata"`
PricingVersion int `json:"pricingVersion"`
Status string `json:"status"`
}
type baseModelScanner interface {
Scan(dest ...any) error
}
func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
rows, err := s.pool.Query(ctx, `
SELECT `+baseModelColumns+`
FROM base_model_catalog
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]BaseModel, 0)
for rows.Next() {
item, err := scanBaseModel(rows)
if err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (BaseModel, error) {
input = normalizeBaseModelInput(input)
capabilities, _ := json.Marshal(emptyObjectIfNil(input.Capabilities))
billingConfig, _ := json.Marshal(emptyObjectIfNil(input.BaseBillingConfig))
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
return scanBaseModel(s.pool.QueryRow(ctx, `
INSERT INTO base_model_catalog (
provider_id, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy, metadata, pricing_version, status
)
VALUES (
(SELECT id FROM model_catalog_providers WHERE provider_key = $1 OR provider_code = $1 LIMIT 1),
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11
)
RETURNING `+baseModelColumns,
input.ProviderKey,
input.CanonicalModelKey,
input.ProviderModelName,
input.ModelType,
input.DisplayName,
capabilities,
billingConfig,
rateLimitPolicy,
metadata,
input.PricingVersion,
input.Status,
))
}
func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelInput) (BaseModel, error) {
input = normalizeBaseModelInput(input)
capabilities, _ := json.Marshal(emptyObjectIfNil(input.Capabilities))
billingConfig, _ := json.Marshal(emptyObjectIfNil(input.BaseBillingConfig))
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
return scanBaseModel(s.pool.QueryRow(ctx, `
UPDATE base_model_catalog
SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = $2 OR provider_code = $2 LIMIT 1),
provider_key = $2,
canonical_model_key = $3,
provider_model_name = $4,
model_type = $5,
display_name = $6,
capabilities = $7,
base_billing_config = $8,
default_rate_limit_policy = $9,
metadata = $10,
pricing_version = $11,
status = $12,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+baseModelColumns,
id,
input.ProviderKey,
input.CanonicalModelKey,
input.ProviderModelName,
input.ModelType,
input.DisplayName,
capabilities,
billingConfig,
rateLimitPolicy,
metadata,
input.PricingVersion,
input.Status,
))
}
func (s *Store) DeleteBaseModel(ctx context.Context, id string) error {
result, err := s.pool.Exec(ctx, `DELETE FROM base_model_catalog WHERE id = $1::uuid`, id)
if err != nil {
return err
}
if result.RowsAffected() == 0 {
return pgx.ErrNoRows
}
return nil
}
func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
var item BaseModel
var capabilities []byte
var billingConfig []byte
var rateLimitPolicy []byte
var metadata []byte
if err := scanner.Scan(
&item.ID,
&item.ProviderKey,
&item.CanonicalModelKey,
&item.ProviderModelName,
&item.ModelType,
&item.DisplayName,
&capabilities,
&billingConfig,
&rateLimitPolicy,
&metadata,
&item.PricingVersion,
&item.Status,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
return BaseModel{}, err
}
item.Capabilities = decodeObject(capabilities)
item.BaseBillingConfig = decodeObject(billingConfig)
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
item.Metadata = decodeObject(metadata)
return item, nil
}
func normalizeBaseModelInput(input BaseModelInput) BaseModelInput {
input.ProviderKey = strings.TrimSpace(input.ProviderKey)
input.CanonicalModelKey = strings.TrimSpace(input.CanonicalModelKey)
input.ProviderModelName = strings.TrimSpace(input.ProviderModelName)
input.ModelType = strings.TrimSpace(input.ModelType)
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.Status = strings.TrimSpace(input.Status)
if input.CanonicalModelKey == "" && input.ProviderKey != "" && input.ProviderModelName != "" {
input.CanonicalModelKey = input.ProviderKey + ":" + input.ProviderModelName
}
if input.DisplayName == "" {
input.DisplayName = input.ProviderModelName
}
if input.ModelType == "" {
input.ModelType = "text_generate"
}
if input.PricingVersion <= 0 {
input.PricingVersion = 1
}
if input.Status == "" {
input.Status = "active"
}
return input
}
apps/api/internal/store/candidates.go
+120
View File
@@ -0,0 +1,120 @@
package store
import (
"context"
"fmt"
)
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string) ([]RuntimeModelCandidate, error) {
rows, err := s.pool.Query(ctx, `
SELECT p.id::text, p.platform_key, p.name, p.provider, COALESCE(p.base_url, ''),
p.auth_type, p.credentials, p.config, p.default_pricing_mode,
p.default_discount_factor::float8, COALESCE(p.pricing_rule_set_id::text, ''),
p.retry_policy, p.rate_limit_policy,
COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
COALESCE(b.provider_model_name, ''), m.model_name, COALESCE(m.model_alias, ''),
m.model_type, m.display_name, m.capabilities, m.capability_override,
COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
m.retry_policy, m.rate_limit_policy
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN runtime_client_states s
ON s.client_id = p.platform_key || ':' || m.model_type || ':' || m.model_name
WHERE p.status = 'enabled'
AND p.deleted_at IS NULL
AND m.enabled = true
AND m.model_type = $2
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
AND (
m.model_name = $1
OR m.model_alias = $1
OR b.canonical_model_key = $1
OR b.provider_model_name = $1
)
ORDER BY effective_priority ASC,
COALESCE(s.limiter_ratio, 0) ASC,
COALESCE(s.running_count, 0) ASC,
COALESCE(s.waiting_count, 0) ASC,
COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
m.created_at ASC`, model, modelType)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]RuntimeModelCandidate, 0)
for rows.Next() {
var item RuntimeModelCandidate
var credentials []byte
var platformConfig []byte
var platformRetryPolicy []byte
var platformRateLimitPolicy []byte
var capabilities []byte
var capabilityOverride []byte
var baseBilling []byte
var billing []byte
var billingOverride []byte
var modelRetryPolicy []byte
var modelRateLimitPolicy []byte
if err := rows.Scan(
&item.PlatformID,
&item.PlatformKey,
&item.PlatformName,
&item.Provider,
&item.BaseURL,
&item.AuthType,
&credentials,
&platformConfig,
&item.DefaultPricingMode,
&item.DefaultDiscountFactor,
&item.PlatformPricingRuleSetID,
&platformRetryPolicy,
&platformRateLimitPolicy,
&item.PlatformPriority,
&item.PlatformModelID,
&item.BaseModelID,
&item.CanonicalModelKey,
&item.ProviderModelName,
&item.ModelName,
&item.ModelAlias,
&item.ModelType,
&item.DisplayName,
&capabilities,
&capabilityOverride,
&baseBilling,
&billing,
&billingOverride,
&item.PricingMode,
&item.DiscountFactor,
&item.ModelPricingRuleSetID,
&modelRetryPolicy,
&modelRateLimitPolicy,
); err != nil {
return nil, err
}
item.Credentials = decodeObject(credentials)
item.PlatformConfig = decodeObject(platformConfig)
item.PlatformRetryPolicy = decodeObject(platformRetryPolicy)
item.PlatformRateLimitPolicy = decodeObject(platformRateLimitPolicy)
item.Capabilities = decodeObject(capabilities)
item.CapabilityOverride = decodeObject(capabilityOverride)
item.BaseBillingConfig = decodeObject(baseBilling)
item.BillingConfig = decodeObject(billing)
item.BillingConfigOverride = decodeObject(billingOverride)
item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, item.ModelName)
item.QueueKey = item.ClientID
items = append(items, item)
}
if err := rows.Err(); err != nil {
return nil, err
}
if len(items) == 0 {
return nil, ErrNoModelCandidate
}
return items, nil
}
apps/api/internal/store/catalog_providers.go
+155
View File
@@ -0,0 +1,155 @@
package store
import (
"context"
"encoding/json"
"strings"
"github.com/jackc/pgx/v5"
)
const catalogProviderColumns = `
id::text, provider_key, COALESCE(NULLIF(provider_code, ''), provider_key) AS provider_code,
display_name, provider_type, COALESCE(icon_path, '') AS icon_path,
COALESCE(source, '') AS source, capability_schema, default_rate_limit_policy,
metadata, status, created_at, updated_at`
type CatalogProviderInput struct {
ProviderKey string `json:"providerKey"`
Code string `json:"code"`
DisplayName string `json:"displayName"`
ProviderType string `json:"providerType"`
IconPath string `json:"iconPath"`
Source string `json:"source"`
CapabilitySchema map[string]any `json:"capabilitySchema"`
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy"`
Metadata map[string]any `json:"metadata"`
Status string `json:"status"`
}
type catalogProviderScanner interface {
Scan(dest ...any) error
}
func (s *Store) CreateCatalogProvider(ctx context.Context, input CatalogProviderInput) (CatalogProvider, error) {
input = normalizeCatalogProviderInput(input)
capabilitySchema, _ := json.Marshal(emptyObjectIfNil(input.CapabilitySchema))
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
return scanCatalogProvider(s.pool.QueryRow(ctx, `
INSERT INTO model_catalog_providers (
provider_key, provider_code, display_name, provider_type, icon_path, source,
capability_schema, default_rate_limit_policy, metadata, status
)
VALUES ($1, $2, $3, $4, NULLIF($5, ''), $6, $7, $8, $9, $10)
RETURNING `+catalogProviderColumns,
input.ProviderKey,
input.Code,
input.DisplayName,
input.ProviderType,
input.IconPath,
input.Source,
capabilitySchema,
rateLimitPolicy,
metadata,
input.Status,
))
}
func (s *Store) UpdateCatalogProvider(ctx context.Context, id string, input CatalogProviderInput) (CatalogProvider, error) {
input = normalizeCatalogProviderInput(input)
capabilitySchema, _ := json.Marshal(emptyObjectIfNil(input.CapabilitySchema))
rateLimitPolicy, _ := json.Marshal(emptyObjectIfNil(input.DefaultRateLimitPolicy))
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
return scanCatalogProvider(s.pool.QueryRow(ctx, `
UPDATE model_catalog_providers
SET provider_key = $2,
provider_code = $3,
display_name = $4,
provider_type = $5,
icon_path = NULLIF($6, ''),
source = $7,
capability_schema = $8,
default_rate_limit_policy = $9,
metadata = $10,
status = $11,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+catalogProviderColumns,
id,
input.ProviderKey,
input.Code,
input.DisplayName,
input.ProviderType,
input.IconPath,
input.Source,
capabilitySchema,
rateLimitPolicy,
metadata,
input.Status,
))
}
func (s *Store) DeleteCatalogProvider(ctx context.Context, id string) error {
result, err := s.pool.Exec(ctx, `DELETE FROM model_catalog_providers WHERE id = $1::uuid`, id)
if err != nil {
return err
}
if result.RowsAffected() == 0 {
return pgx.ErrNoRows
}
return nil
}
func scanCatalogProvider(scanner catalogProviderScanner) (CatalogProvider, error) {
var item CatalogProvider
var capabilitySchema []byte
var rateLimitPolicy []byte
var metadata []byte
if err := scanner.Scan(
&item.ID,
&item.ProviderKey,
&item.Code,
&item.DisplayName,
&item.ProviderType,
&item.IconPath,
&item.Source,
&capabilitySchema,
&rateLimitPolicy,
&metadata,
&item.Status,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
return CatalogProvider{}, err
}
item.CapabilitySchema = decodeObject(capabilitySchema)
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
item.Metadata = decodeObject(metadata)
return item, nil
}
func normalizeCatalogProviderInput(input CatalogProviderInput) CatalogProviderInput {
input.ProviderKey = strings.TrimSpace(input.ProviderKey)
input.Code = strings.TrimSpace(input.Code)
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.ProviderType = strings.TrimSpace(input.ProviderType)
input.IconPath = strings.TrimSpace(input.IconPath)
input.Source = strings.TrimSpace(input.Source)
input.Status = strings.TrimSpace(input.Status)
if input.Code == "" {
input.Code = input.ProviderKey
}
if input.ProviderType == "" {
input.ProviderType = "openai"
}
if input.Source == "" {
input.Source = "gateway"
}
if input.Status == "" {
input.Status = "active"
}
return input
}
apps/api/internal/store/platform_models.go
+206
View File
@@ -0,0 +1,206 @@
package store
import (
"context"
"encoding/json"
"strings"
"github.com/jackc/pgx/v5"
)
type modelCatalogSnapshot struct {
ID string
ProviderKey string
CanonicalModelKey string
ProviderModelName string
ModelType string
DisplayName string
Capabilities map[string]any
BaseBillingConfig map[string]any
DefaultRateLimitPolicy map[string]any
}
func (s *Store) CreatePlatformModel(ctx context.Context, input CreatePlatformModelInput) (PlatformModel, error) {
base, err := s.lookupBaseModel(ctx, input.BaseModelID, input.CanonicalModelKey, input.ModelName)
if err != nil && !IsNotFound(err) {
return PlatformModel{}, err
}
if input.ModelType == "" {
input.ModelType = base.ModelType
}
if input.ModelName == "" {
input.ModelName = base.ProviderModelName
}
if input.DisplayName == "" {
input.DisplayName = firstNonEmpty(base.DisplayName, input.ModelName)
}
if input.PricingMode == "" {
input.PricingMode = "inherit_discount"
}
capabilities := input.Capabilities
if len(capabilities) == 0 {
capabilities = mergeObjects(base.Capabilities, input.CapabilityOverride)
}
billingConfig := input.BillingConfig
if len(billingConfig) == 0 {
billingConfig = mergeObjects(base.BaseBillingConfig, input.BillingConfigOverride)
}
rateLimitPolicy := input.RateLimitPolicy
if len(rateLimitPolicy) == 0 {
rateLimitPolicy = base.DefaultRateLimitPolicy
}
capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
billingOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingConfigOverride))
billingJSON, _ := json.Marshal(emptyObjectIfNil(billingConfig))
permissionJSON, _ := json.Marshal(emptyObjectIfNil(input.PermissionConfig))
retryJSON, _ := json.Marshal(emptyObjectIfNil(input.RetryPolicy))
rateLimitJSON, _ := json.Marshal(emptyObjectIfNil(rateLimitPolicy))
discount := any(nil)
if input.DiscountFactor > 0 {
discount = input.DiscountFactor
}
baseID := any(nil)
if base.ID != "" {
baseID = base.ID
}
var model PlatformModel
var capabilityOverrideBytes []byte
var capabilitiesBytes []byte
var billingOverrideBytes []byte
var billingBytes []byte
err = s.pool.QueryRow(ctx, `
INSERT INTO platform_models (
platform_id, base_model_id, model_name, model_alias, model_type, display_name,
capability_override, capabilities, pricing_mode, discount_factor,
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy, enabled
)
VALUES (
$1::uuid, $2::uuid, $3, NULLIF($4, ''), $5, $6,
$7::jsonb, $8::jsonb, $9, $10::numeric,
NULLIF($11, '')::uuid, $12::jsonb, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, true
)
ON CONFLICT (platform_id, model_name, model_type) DO UPDATE
SET base_model_id = EXCLUDED.base_model_id,
model_alias = EXCLUDED.model_alias,
display_name = EXCLUDED.display_name,
capability_override = EXCLUDED.capability_override,
capabilities = EXCLUDED.capabilities,
pricing_mode = EXCLUDED.pricing_mode,
discount_factor = EXCLUDED.discount_factor,
pricing_rule_set_id = EXCLUDED.pricing_rule_set_id,
billing_config_override = EXCLUDED.billing_config_override,
billing_config = EXCLUDED.billing_config,
permission_config = EXCLUDED.permission_config,
retry_policy = EXCLUDED.retry_policy,
rate_limit_policy = EXCLUDED.rate_limit_policy,
enabled = true,
updated_at = now()
RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_name,
COALESCE(model_alias, ''), model_type, display_name, capability_override,
capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config, enabled, created_at, updated_at`,
input.PlatformID,
baseID,
input.ModelName,
input.ModelAlias,
input.ModelType,
input.DisplayName,
string(capabilityOverrideJSON),
string(capabilitiesJSON),
input.PricingMode,
discount,
input.PricingRuleSetID,
string(billingOverrideJSON),
string(billingJSON),
string(permissionJSON),
string(retryJSON),
string(rateLimitJSON),
).Scan(
&model.ID,
&model.PlatformID,
&model.BaseModelID,
&model.ModelName,
&model.ModelAlias,
&model.ModelType,
&model.DisplayName,
&capabilityOverrideBytes,
&capabilitiesBytes,
&model.PricingMode,
&model.DiscountFactor,
&model.PricingRuleSetID,
&billingOverrideBytes,
&billingBytes,
&model.Enabled,
&model.CreatedAt,
&model.UpdatedAt,
)
if err != nil {
return PlatformModel{}, err
}
model.CapabilityOverride = decodeObject(capabilityOverrideBytes)
model.Capabilities = decodeObject(capabilitiesBytes)
model.BillingConfigOverride = decodeObject(billingOverrideBytes)
model.BillingConfig = decodeObject(billingBytes)
return model, nil
}
func (s *Store) lookupBaseModel(ctx context.Context, id string, canonicalKey string, modelName string) (modelCatalogSnapshot, error) {
var item modelCatalogSnapshot
var capabilities []byte
var billingConfig []byte
var rateLimitPolicy []byte
err := s.pool.QueryRow(ctx, `
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy
FROM base_model_catalog
WHERE ($1 <> '' AND id = NULLIF($1, '')::uuid)
OR ($2 <> '' AND canonical_model_key = $2)
OR ($3 <> '' AND provider_model_name = $3)
ORDER BY CASE WHEN id::text = $1 THEN 0 WHEN canonical_model_key = $2 THEN 1 ELSE 2 END
LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSpace(modelName)).Scan(
&item.ID,
&item.ProviderKey,
&item.CanonicalModelKey,
&item.ProviderModelName,
&item.ModelType,
&item.DisplayName,
&capabilities,
&billingConfig,
&rateLimitPolicy,
)
if err != nil {
if err == pgx.ErrNoRows {
return modelCatalogSnapshot{}, err
}
return modelCatalogSnapshot{}, err
}
item.Capabilities = decodeObject(capabilities)
item.BaseBillingConfig = decodeObject(billingConfig)
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
return item, nil
}
func mergeObjects(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
}
if len(out) == 0 {
return nil
}
return out
}
func emptyObjectIfNil(value map[string]any) map[string]any {
if value == nil {
return map[string]any{}
}
return value
}
apps/api/internal/store/postgres.go
+95 -121
View File
@@ -60,6 +60,7 @@ type Platform struct {
Priority int `json:"priority"`
DefaultPricingMode string `json:"defaultPricingMode"`
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
Config map[string]any `json:"config,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
@@ -75,6 +76,7 @@ type CreatePlatformInput struct {
Config map[string]any `json:"config"`
DefaultPricingMode string `json:"defaultPricingMode"`
DefaultDiscountFactor float64 `json:"defaultDiscountFactor"`
PricingRuleSetID string `json:"pricingRuleSetId"`
Priority int `json:"priority"`
}
@@ -123,6 +125,7 @@ type PlatformModel struct {
Capabilities map[string]any `json:"capabilities,omitempty"`
PricingMode string `json:"pricingMode"`
DiscountFactor float64 `json:"discountFactor,omitempty"`
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
BillingConfigOverride map[string]any `json:"billingConfigOverride,omitempty"`
BillingConfig map[string]any `json:"billingConfig,omitempty"`
Enabled bool `json:"enabled"`
@@ -133,10 +136,14 @@ type PlatformModel struct {
type CatalogProvider struct {
ID string `json:"id"`
ProviderKey string `json:"providerKey"`
Code string `json:"code"`
DisplayName string `json:"displayName"`
ProviderType string `json:"providerType"`
IconPath string `json:"iconPath,omitempty"`
Source string `json:"source,omitempty"`
CapabilitySchema map[string]any `json:"capabilitySchema,omitempty"`
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
Status string `json:"status"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
@@ -152,6 +159,7 @@ type BaseModel struct {
Capabilities map[string]any `json:"capabilities,omitempty"`
BaseBillingConfig map[string]any `json:"baseBillingConfig,omitempty"`
DefaultRateLimitPolicy map[string]any `json:"defaultRateLimitPolicy,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
PricingVersion int `json:"pricingVersion"`
Status string `json:"status"`
CreatedAt time.Time `json:"createdAt"`
@@ -159,17 +167,40 @@ type BaseModel struct {
}
type PricingRule struct {
ID string `json:"id"`
ScopeType string `json:"scopeType"`
ScopeID string `json:"scopeId,omitempty"`
ResourceType string `json:"resourceType"`
Unit string `json:"unit"`
BasePrice float64 `json:"basePrice"`
Currency string `json:"currency"`
BaseWeight map[string]any `json:"baseWeight,omitempty"`
DynamicWeight map[string]any `json:"dynamicWeight,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
ID string `json:"id"`
RuleSetID string `json:"ruleSetId,omitempty"`
RuleKey string `json:"ruleKey"`
DisplayName string `json:"displayName"`
ScopeType string `json:"scopeType"`
ScopeID string `json:"scopeId,omitempty"`
ResourceType string `json:"resourceType"`
Unit string `json:"unit"`
BasePrice float64 `json:"basePrice"`
Currency string `json:"currency"`
BaseWeight map[string]any `json:"baseWeight,omitempty"`
DynamicWeight map[string]any `json:"dynamicWeight,omitempty"`
CalculatorType string `json:"calculatorType"`
DimensionSchema map[string]any `json:"dimensionSchema,omitempty"`
FormulaConfig map[string]any `json:"formulaConfig,omitempty"`
Priority int `json:"priority"`
Status string `json:"status"`
Metadata map[string]any `json:"metadata,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type PricingRuleSet struct {
ID string `json:"id"`
RuleSetKey string `json:"ruleSetKey"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Category string `json:"category"`
Currency string `json:"currency"`
Status string `json:"status"`
Metadata map[string]any `json:"metadata,omitempty"`
Rules []PricingRule `json:"rules,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type GatewayTenant struct {
@@ -305,7 +336,8 @@ type TaskEvent struct {
func (s *Store) ListPlatforms(ctx context.Context) ([]Platform, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
default_pricing_mode, default_discount_factor::float8, config, created_at, updated_at
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''),
config, created_at, updated_at
FROM integration_platforms
ORDER BY priority ASC, created_at DESC`)
if err != nil {
@@ -328,6 +360,7 @@ ORDER BY priority ASC, created_at DESC`)
&platform.Priority,
&platform.DefaultPricingMode,
&platform.DefaultDiscountFactor,
&platform.PricingRuleSetID,
&configBytes,
&platform.CreatedAt,
&platform.UpdatedAt,
@@ -355,11 +388,11 @@ func (s *Store) CreatePlatform(ctx context.Context, input CreatePlatformInput) (
var platform Platform
var configBytes []byte
err := s.pool.QueryRow(ctx, `
INSERT INTO integration_platforms (provider, platform_key, name, base_url, auth_type, credentials, config, default_pricing_mode, default_discount_factor, priority)
VALUES ($1, COALESCE(NULLIF($2, ''), 'platform_' || replace(gen_random_uuid()::text, '-', '')), $3, $4, $5, $6, $7, $8, $9, $10)
INSERT INTO integration_platforms (provider, platform_key, name, base_url, auth_type, credentials, config, default_pricing_mode, default_discount_factor, pricing_rule_set_id, priority)
VALUES ($1, COALESCE(NULLIF($2, ''), 'platform_' || replace(gen_random_uuid()::text, '-', '')), $3, $4, $5, $6, $7, $8, $9, NULLIF($10, '')::uuid, $11)
RETURNING id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_type, status, priority,
default_pricing_mode, default_discount_factor::float8, config, created_at, updated_at`,
input.Provider, input.PlatformKey, input.Name, input.BaseURL, input.AuthType, credentials, config, input.DefaultPricingMode, input.DefaultDiscountFactor, input.Priority,
default_pricing_mode, default_discount_factor::float8, COALESCE(pricing_rule_set_id::text, ''), config, created_at, updated_at`,
input.Provider, input.PlatformKey, input.Name, input.BaseURL, input.AuthType, credentials, config, input.DefaultPricingMode, input.DefaultDiscountFactor, input.PricingRuleSetID, input.Priority,
).Scan(
&platform.ID,
&platform.Provider,
@@ -371,6 +404,7 @@ RETURNING id::text, provider, platform_key, name, COALESCE(base_url, ''), auth_t
&platform.Priority,
&platform.DefaultPricingMode,
&platform.DefaultDiscountFactor,
&platform.PricingRuleSetID,
&configBytes,
&platform.CreatedAt,
&platform.UpdatedAt,
@@ -387,7 +421,8 @@ func (s *Store) ListModels(ctx context.Context) ([]PlatformModel, error) {
SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name,
m.model_name, COALESCE(m.model_alias, ''), m.model_type, m.display_name,
m.capability_override, m.capabilities, m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
m.billing_config_override, m.billing_config, m.enabled, m.created_at, m.updated_at
COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
m.enabled, m.created_at, m.updated_at
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
ORDER BY m.model_type ASC, m.model_name ASC`)
@@ -417,6 +452,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`)
&capabilities,
&model.PricingMode,
&model.DiscountFactor,
&model.PricingRuleSetID,
&billingConfigOverride,
&billingConfig,
&model.Enabled,
@@ -436,8 +472,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`)
func (s *Store) ListCatalogProviders(ctx context.Context) ([]CatalogProvider, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, provider_key, display_name, provider_type, capability_schema,
default_rate_limit_policy, status, created_at, updated_at
SELECT `+catalogProviderColumns+`
FROM model_catalog_providers
ORDER BY provider_key ASC`)
if err != nil {
@@ -447,67 +482,10 @@ ORDER BY provider_key ASC`)
items := make([]CatalogProvider, 0)
for rows.Next() {
var item CatalogProvider
var capabilitySchema []byte
var rateLimitPolicy []byte
if err := rows.Scan(
&item.ID,
&item.ProviderKey,
&item.DisplayName,
&item.ProviderType,
&capabilitySchema,
&rateLimitPolicy,
&item.Status,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
item, err := scanCatalogProvider(rows)
if err != nil {
return nil, err
}
item.CapabilitySchema = decodeObject(capabilitySchema)
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
items = append(items, item)
}
return items, rows.Err()
}
func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy, pricing_version,
status, created_at, updated_at
FROM base_model_catalog
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]BaseModel, 0)
for rows.Next() {
var item BaseModel
var capabilities []byte
var billingConfig []byte
var rateLimitPolicy []byte
if err := rows.Scan(
&item.ID,
&item.ProviderKey,
&item.CanonicalModelKey,
&item.ProviderModelName,
&item.ModelType,
&item.DisplayName,
&capabilities,
&billingConfig,
&rateLimitPolicy,
&item.PricingVersion,
&item.Status,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
return nil, err
}
item.Capabilities = decodeObject(capabilities)
item.BaseBillingConfig = decodeObject(billingConfig)
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
items = append(items, item)
}
return items, rows.Err()
@@ -516,9 +494,11 @@ ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
func (s *Store) ListPricingRules(ctx context.Context) ([]PricingRule, error) {
rows, err := s.pool.Query(ctx, `
SELECT id::text, scope_type, COALESCE(scope_id::text, ''), resource_type, unit,
base_price::float8, currency, base_weight, dynamic_weight, created_at, updated_at
base_price::float8, currency, base_weight, dynamic_weight,
COALESCE(rule_set_id::text, ''), rule_key, display_name, calculator_type,
dimension_schema, formula_config, priority, status, metadata, created_at, updated_at
FROM model_pricing_rules
ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
ORDER BY COALESCE(rule_set_id::text, ''), priority ASC, resource_type ASC, created_at DESC`)
if err != nil {
return nil, err
}
@@ -529,6 +509,9 @@ ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
var item PricingRule
var baseWeight []byte
var dynamicWeight []byte
var dimensionSchema []byte
var formulaConfig []byte
var metadata []byte
if err := rows.Scan(
&item.ID,
&item.ScopeType,
@@ -539,6 +522,15 @@ ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
&item.Currency,
&baseWeight,
&dynamicWeight,
&item.RuleSetID,
&item.RuleKey,
&item.DisplayName,
&item.CalculatorType,
&dimensionSchema,
&formulaConfig,
&item.Priority,
&item.Status,
&metadata,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
@@ -546,6 +538,9 @@ ORDER BY scope_type ASC, resource_type ASC, created_at DESC`)
}
item.BaseWeight = decodeObject(baseWeight)
item.DynamicWeight = decodeObject(dynamicWeight)
item.DimensionSchema = decodeObject(dimensionSchema)
item.FormulaConfig = decodeObject(formulaConfig)
item.Metadata = decodeObject(metadata)
items = append(items, item)
}
return items, rows.Err()
@@ -969,6 +964,17 @@ SELECT NOT EXISTS (
).Scan(&tenantID); err != nil {
return GatewayUser{}, err
}
_ = tx.QueryRow(ctx, `
SELECT COALESCE(default_user_group_id::text, '')
FROM gateway_tenants
WHERE id = $1::uuid`, tenantID).Scan(&userGroupID)
if userGroupID == "" {
_ = tx.QueryRow(ctx, `
SELECT id::text
FROM gateway_user_groups
WHERE group_key = 'default' AND status = 'active'
LIMIT 1`).Scan(&userGroupID)
}
if invitationCode != "" {
if err := tx.QueryRow(ctx, `
SELECT i.id::text, COALESCE(i.created_by::text, '')
@@ -1182,17 +1188,8 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
requestBody, _ := json.Marshal(input.Request)
runMode := normalizeRunMode(input.RunMode, input.Request)
status := "queued"
result := map[string]any(nil)
billings := []any(nil)
finished := false
if runMode == "simulation" {
status = "succeeded"
result = simulationResult(input.Kind, input.Model)
billings = simulationBillings(input.Kind, input.Model)
finished = true
}
resultBody, _ := json.Marshal(result)
billingsBody, _ := json.Marshal(billings)
resultBody, _ := json.Marshal(map[string]any(nil))
billingsBody, _ := json.Marshal([]any(nil))
tx, err := s.pool.Begin(ctx)
if err != nil {
@@ -1213,12 +1210,12 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
RETURNING id::text, kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model, request, status, result, billings, COALESCE(error, ''), created_at, updated_at`,
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, finished,
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, false,
).Scan(&task.ID, &task.Kind, &task.RunMode, &task.UserID, &task.GatewayUserID, &task.UserSource, &task.GatewayTenantID, &task.TenantID, &task.TenantKey, &task.UserGroupID, &task.UserGroupKey, &task.Model, &requestBytes, &task.Status, &resultBytes, &billingsBytes, &task.Error, &task.CreatedAt, &task.UpdatedAt)
if err != nil {
return GatewayTask{}, err
}
events := taskEventsForCreate(task.ID, runMode, status, result)
events := taskEventsForCreate(task.ID, runMode, status, nil)
for _, event := range events {
payload, _ := json.Marshal(event.Payload)
if _, err := tx.Exec(ctx, `
@@ -1300,6 +1297,10 @@ func IsNotFound(err error) bool {
return err == pgx.ErrNoRows
}
func IsUniqueViolation(err error) bool {
return isUniqueViolation(err)
}
func isUniqueViolation(err error) bool {
var pgErr *pgconn.PgError
return errors.As(err, &pgErr) && pgErr.Code == "23505"
@@ -1452,7 +1453,7 @@ func simulationBillings(kind string, model string) []any {
}
func taskEventsForCreate(taskID string, runMode string, status string, result map[string]any) []TaskEvent {
events := []TaskEvent{{
return []TaskEvent{{
TaskID: taskID,
Seq: 1,
EventType: "task.accepted",
@@ -1463,33 +1464,6 @@ func taskEventsForCreate(taskID string, runMode string, status string, result ma
Payload: map[string]any{"taskId": taskID},
Simulated: runMode == "simulation",
}}
if runMode != "simulation" {
return events
}
return append(events,
TaskEvent{
TaskID: taskID,
Seq: 2,
EventType: "task.progress",
Status: "running",
Phase: "simulation",
Progress: 0.5,
Message: "simulation client running",
Payload: map[string]any{"taskId": taskID},
Simulated: true,
},
TaskEvent{
TaskID: taskID,
Seq: 3,
EventType: "task.completed",
Status: status,
Phase: "completed",
Progress: 1,
Message: "simulation completed",
Payload: map[string]any{"taskId": taskID, "result": result},
Simulated: true,
},
)
}
func decodeObject(bytes []byte) map[string]any {
apps/api/internal/store/pricing_rules.go
+334
View File
@@ -0,0 +1,334 @@
package store
import (
"context"
"encoding/json"
"strings"
"github.com/jackc/pgx/v5"
)
const pricingRuleSetColumns = `
id::text, rule_set_key, name, COALESCE(description, ''), category, currency,
status, metadata, created_at, updated_at`
const pricingRuleColumns = `
id::text, COALESCE(rule_set_id::text, ''), rule_key, display_name, scope_type,
COALESCE(scope_id::text, ''), resource_type, unit, base_price::float8, currency,
base_weight, dynamic_weight, calculator_type, dimension_schema, formula_config,
priority, status, metadata, created_at, updated_at`
type PricingRuleInput struct {
RuleKey string `json:"ruleKey"`
DisplayName string `json:"displayName"`
ResourceType string `json:"resourceType"`
Unit string `json:"unit"`
BasePrice float64 `json:"basePrice"`
Currency string `json:"currency"`
BaseWeight map[string]any `json:"baseWeight"`
DynamicWeight map[string]any `json:"dynamicWeight"`
CalculatorType string `json:"calculatorType"`
DimensionSchema map[string]any `json:"dimensionSchema"`
FormulaConfig map[string]any `json:"formulaConfig"`
Priority int `json:"priority"`
Status string `json:"status"`
Metadata map[string]any `json:"metadata"`
}
type PricingRuleSetInput struct {
RuleSetKey string `json:"ruleSetKey"`
Name string `json:"name"`
Description string `json:"description"`
Category string `json:"category"`
Currency string `json:"currency"`
Status string `json:"status"`
Metadata map[string]any `json:"metadata"`
Rules []PricingRuleInput `json:"rules"`
}
type pricingScanner interface {
Scan(dest ...any) error
}
func (s *Store) ListPricingRuleSets(ctx context.Context) ([]PricingRuleSet, error) {
rows, err := s.pool.Query(ctx, `SELECT `+pricingRuleSetColumns+` FROM model_pricing_rule_sets ORDER BY category ASC, name ASC`)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]PricingRuleSet, 0)
byID := map[string]int{}
for rows.Next() {
item, err := scanPricingRuleSet(rows)
if err != nil {
return nil, err
}
byID[item.ID] = len(items)
items = append(items, item)
}
if err := rows.Err(); err != nil {
return nil, err
}
ruleRows, err := s.pool.Query(ctx, `
SELECT `+pricingRuleColumns+`
FROM model_pricing_rules
WHERE rule_set_id IS NOT NULL
ORDER BY rule_set_id, priority ASC, resource_type ASC, rule_key ASC`)
if err != nil {
return nil, err
}
defer ruleRows.Close()
for ruleRows.Next() {
rule, err := scanPricingRule(ruleRows)
if err != nil {
return nil, err
}
if index, ok := byID[rule.RuleSetID]; ok {
items[index].Rules = append(items[index].Rules, rule)
}
}
return items, ruleRows.Err()
}
func (s *Store) CreatePricingRuleSet(ctx context.Context, input PricingRuleSetInput) (PricingRuleSet, error) {
input = normalizePricingRuleSet(input)
tx, err := s.pool.Begin(ctx)
if err != nil {
return PricingRuleSet{}, err
}
defer tx.Rollback(ctx)
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
item, err := scanPricingRuleSet(tx.QueryRow(ctx, `
INSERT INTO model_pricing_rule_sets (rule_set_key, name, description, category, currency, status, metadata)
VALUES ($1, $2, NULLIF($3, ''), $4, $5, $6, $7)
RETURNING `+pricingRuleSetColumns,
input.RuleSetKey, input.Name, input.Description, input.Category, input.Currency, input.Status, metadata,
))
if err != nil {
return PricingRuleSet{}, err
}
if err := insertPricingRules(ctx, tx, item.ID, input.Currency, input.Rules); err != nil {
return PricingRuleSet{}, err
}
if err := tx.Commit(ctx); err != nil {
return PricingRuleSet{}, err
}
item.Rules = pricingInputsToRules(item.ID, input.Currency, input.Rules)
return item, nil
}
func (s *Store) UpdatePricingRuleSet(ctx context.Context, id string, input PricingRuleSetInput) (PricingRuleSet, error) {
input = normalizePricingRuleSet(input)
tx, err := s.pool.Begin(ctx)
if err != nil {
return PricingRuleSet{}, err
}
defer tx.Rollback(ctx)
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
item, err := scanPricingRuleSet(tx.QueryRow(ctx, `
UPDATE model_pricing_rule_sets
SET rule_set_key = $2,
name = $3,
description = NULLIF($4, ''),
category = $5,
currency = $6,
status = $7,
metadata = $8,
updated_at = now()
WHERE id = $1::uuid
RETURNING `+pricingRuleSetColumns,
id, input.RuleSetKey, input.Name, input.Description, input.Category, input.Currency, input.Status, metadata,
))
if err != nil {
return PricingRuleSet{}, err
}
if _, err := tx.Exec(ctx, `DELETE FROM model_pricing_rules WHERE rule_set_id = $1::uuid`, id); err != nil {
return PricingRuleSet{}, err
}
if err := insertPricingRules(ctx, tx, item.ID, input.Currency, input.Rules); err != nil {
return PricingRuleSet{}, err
}
if err := tx.Commit(ctx); err != nil {
return PricingRuleSet{}, err
}
item.Rules = pricingInputsToRules(item.ID, input.Currency, input.Rules)
return item, nil
}
func (s *Store) DeletePricingRuleSet(ctx context.Context, id string) error {
result, err := s.pool.Exec(ctx, `DELETE FROM model_pricing_rule_sets WHERE id = $1::uuid`, id)
if err != nil {
return err
}
if result.RowsAffected() == 0 {
return pgx.ErrNoRows
}
return nil
}
func insertPricingRules(ctx context.Context, tx pgx.Tx, ruleSetID string, defaultCurrency string, rules []PricingRuleInput) error {
for index, rule := range rules {
rule = normalizePricingRule(rule, index, defaultCurrency)
baseWeight, _ := json.Marshal(emptyObjectIfNil(rule.BaseWeight))
dynamicWeight, _ := json.Marshal(emptyObjectIfNil(rule.DynamicWeight))
dimensionSchema, _ := json.Marshal(emptyObjectIfNil(rule.DimensionSchema))
formulaConfig, _ := json.Marshal(emptyObjectIfNil(rule.FormulaConfig))
metadata, _ := json.Marshal(emptyObjectIfNil(rule.Metadata))
if _, err := tx.Exec(ctx, `
INSERT INTO model_pricing_rules (
rule_set_id, rule_key, display_name, scope_type, scope_id, resource_type,
unit, base_price, currency, base_weight, dynamic_weight, calculator_type,
dimension_schema, formula_config, priority, status, metadata
)
VALUES ($1::uuid, $2, $3, 'rule_set', $1::uuid, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)`,
ruleSetID, rule.RuleKey, rule.DisplayName, rule.ResourceType, rule.Unit,
rule.BasePrice, rule.Currency, baseWeight, dynamicWeight, rule.CalculatorType,
dimensionSchema, formulaConfig, rule.Priority, rule.Status, metadata,
); err != nil {
return err
}
}
return nil
}
func scanPricingRuleSet(scanner pricingScanner) (PricingRuleSet, error) {
var item PricingRuleSet
var metadata []byte
if err := scanner.Scan(
&item.ID,
&item.RuleSetKey,
&item.Name,
&item.Description,
&item.Category,
&item.Currency,
&item.Status,
&metadata,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
return PricingRuleSet{}, err
}
item.Metadata = decodeObject(metadata)
item.Rules = []PricingRule{}
return item, nil
}
func scanPricingRule(scanner pricingScanner) (PricingRule, error) {
var item PricingRule
var baseWeight []byte
var dynamicWeight []byte
var dimensionSchema []byte
var formulaConfig []byte
var metadata []byte
if err := scanner.Scan(
&item.ID,
&item.RuleSetID,
&item.RuleKey,
&item.DisplayName,
&item.ScopeType,
&item.ScopeID,
&item.ResourceType,
&item.Unit,
&item.BasePrice,
&item.Currency,
&baseWeight,
&dynamicWeight,
&item.CalculatorType,
&dimensionSchema,
&formulaConfig,
&item.Priority,
&item.Status,
&metadata,
&item.CreatedAt,
&item.UpdatedAt,
); err != nil {
return PricingRule{}, err
}
item.BaseWeight = decodeObject(baseWeight)
item.DynamicWeight = decodeObject(dynamicWeight)
item.DimensionSchema = decodeObject(dimensionSchema)
item.FormulaConfig = decodeObject(formulaConfig)
item.Metadata = decodeObject(metadata)
return item, nil
}
func normalizePricingRuleSet(input PricingRuleSetInput) PricingRuleSetInput {
input.RuleSetKey = strings.TrimSpace(input.RuleSetKey)
input.Name = strings.TrimSpace(input.Name)
input.Description = strings.TrimSpace(input.Description)
input.Category = strings.TrimSpace(input.Category)
input.Currency = strings.TrimSpace(input.Currency)
input.Status = strings.TrimSpace(input.Status)
if input.Category == "" {
input.Category = "custom"
}
if input.Currency == "" {
input.Currency = "resource"
}
if input.Status == "" {
input.Status = "active"
}
return input
}
func normalizePricingRule(input PricingRuleInput, index int, defaultCurrency string) PricingRuleInput {
input.RuleKey = strings.TrimSpace(input.RuleKey)
input.DisplayName = strings.TrimSpace(input.DisplayName)
input.ResourceType = strings.TrimSpace(input.ResourceType)
input.Unit = strings.TrimSpace(input.Unit)
input.Currency = strings.TrimSpace(input.Currency)
input.CalculatorType = strings.TrimSpace(input.CalculatorType)
input.Status = strings.TrimSpace(input.Status)
if input.RuleKey == "" {
input.RuleKey = "rule_" + strings.ReplaceAll(input.ResourceType+"_"+input.Unit, " ", "_")
}
if input.DisplayName == "" {
input.DisplayName = input.ResourceType
}
if input.Unit == "" {
input.Unit = "item"
}
if input.Currency == "" {
input.Currency = defaultCurrency
}
if input.CalculatorType == "" {
input.CalculatorType = "unit_weight"
}
if input.Priority == 0 {
input.Priority = (index + 1) * 10
}
if input.Status == "" {
input.Status = "active"
}
return input
}
func pricingInputsToRules(ruleSetID string, defaultCurrency string, rules []PricingRuleInput) []PricingRule {
items := make([]PricingRule, 0, len(rules))
for index, input := range rules {
input = normalizePricingRule(input, index, defaultCurrency)
items = append(items, PricingRule{
RuleSetID: ruleSetID,
RuleKey: input.RuleKey,
DisplayName: input.DisplayName,
ScopeType: "rule_set",
ScopeID: ruleSetID,
ResourceType: input.ResourceType,
Unit: input.Unit,
BasePrice: input.BasePrice,
Currency: input.Currency,
BaseWeight: emptyObjectIfNil(input.BaseWeight),
DynamicWeight: emptyObjectIfNil(input.DynamicWeight),
CalculatorType: input.CalculatorType,
DimensionSchema: emptyObjectIfNil(input.DimensionSchema),
FormulaConfig: emptyObjectIfNil(input.FormulaConfig),
Priority: input.Priority,
Status: input.Status,
Metadata: emptyObjectIfNil(input.Metadata),
})
}
return items
}
apps/api/internal/store/rate_limits.go
+109
View File
@@ -0,0 +1,109 @@
package store
import (
"context"
"errors"
)
func (s *Store) ReserveRateLimits(ctx context.Context, taskID string, reservations []RateLimitReservation) (RateLimitResult, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return RateLimitResult{}, err
}
defer tx.Rollback(ctx)
result := RateLimitResult{}
for _, reservation := range reservations {
if reservation.Limit <= 0 || reservation.Amount <= 0 {
continue
}
if reservation.Metric == "" || reservation.Amount > reservation.Limit {
return RateLimitResult{}, ErrRateLimited
}
if reservation.WindowSeconds <= 0 {
reservation.WindowSeconds = 60
}
if reservation.Metric == "concurrent" {
if reservation.LeaseTTLSeconds <= 0 {
reservation.LeaseTTLSeconds = 120
}
var active float64
if err := tx.QueryRow(ctx, `
SELECT COALESCE(SUM(lease_value), 0)::float8
FROM gateway_concurrency_leases
WHERE scope_type = $1
AND scope_key = $2
AND released_at IS NULL
AND expires_at > now()`,
reservation.ScopeType,
reservation.ScopeKey,
).Scan(&active); err != nil {
return RateLimitResult{}, err
}
if active+reservation.Amount > reservation.Limit {
return RateLimitResult{}, ErrRateLimited
}
var leaseID string
if err := tx.QueryRow(ctx, `
INSERT INTO gateway_concurrency_leases (task_id, scope_type, scope_key, lease_value, expires_at)
VALUES ($1::uuid, $2, $3, $4, now() + ($5::int * interval '1 second'))
RETURNING id::text`,
taskID,
reservation.ScopeType,
reservation.ScopeKey,
reservation.Amount,
reservation.LeaseTTLSeconds,
).Scan(&leaseID); err != nil {
return RateLimitResult{}, err
}
result.LeaseIDs = append(result.LeaseIDs, leaseID)
continue
}
tag, err := tx.Exec(ctx, `
INSERT INTO gateway_rate_limit_counters (
scope_type, scope_key, metric, window_start, limit_value, used_value, reserved_value, reset_at
)
VALUES (
$1, $2, $3, date_trunc('minute', now()), $4, $5, 0,
date_trunc('minute', now()) + ($6::int * interval '1 second')
)
ON CONFLICT (scope_type, scope_key, metric, window_start) DO UPDATE
SET limit_value = EXCLUDED.limit_value,
used_value = gateway_rate_limit_counters.used_value + EXCLUDED.used_value,
reset_at = EXCLUDED.reset_at,
updated_at = now()
WHERE gateway_rate_limit_counters.used_value + EXCLUDED.used_value <= EXCLUDED.limit_value`,
reservation.ScopeType,
reservation.ScopeKey,
reservation.Metric,
reservation.Limit,
reservation.Amount,
reservation.WindowSeconds,
)
if err != nil {
return RateLimitResult{}, err
}
if tag.RowsAffected() == 0 {
return RateLimitResult{}, ErrRateLimited
}
}
return result, tx.Commit(ctx)
}
func (s *Store) ReleaseConcurrencyLeases(ctx context.Context, leaseIDs []string) error {
if len(leaseIDs) == 0 {
return nil
}
for _, leaseID := range leaseIDs {
if leaseID == "" {
continue
}
if _, err := s.pool.Exec(ctx, `
UPDATE gateway_concurrency_leases
SET released_at = now()
WHERE id = $1::uuid AND released_at IS NULL`, leaseID); err != nil && !errors.Is(err, ErrRateLimited) {
return err
}
}
return nil
}
apps/api/internal/store/runtime_types.go
+101
View File
@@ -0,0 +1,101 @@
package store
import "errors"
var (
ErrNoModelCandidate = errors.New("no enabled platform model matches request")
ErrRateLimited = errors.New("rate limit exceeded")
)
type CreatePlatformModelInput struct {
PlatformID string `json:"platformId"`
BaseModelID string `json:"baseModelId"`
CanonicalModelKey string `json:"canonicalModelKey"`
ModelName string `json:"modelName"`
ModelAlias string `json:"modelAlias"`
ModelType string `json:"modelType"`
DisplayName string `json:"displayName"`
CapabilityOverride map[string]any `json:"capabilityOverride"`
Capabilities map[string]any `json:"capabilities"`
PricingMode string `json:"pricingMode"`
DiscountFactor float64 `json:"discountFactor"`
PricingRuleSetID string `json:"pricingRuleSetId"`
BillingConfigOverride map[string]any `json:"billingConfigOverride"`
BillingConfig map[string]any `json:"billingConfig"`
PermissionConfig map[string]any `json:"permissionConfig"`
RetryPolicy map[string]any `json:"retryPolicy"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy"`
Enabled bool `json:"enabled"`
}
type RuntimeModelCandidate struct {
PlatformID string
PlatformKey string
PlatformName string
Provider string
BaseURL string
AuthType string
Credentials map[string]any
PlatformConfig map[string]any
DefaultPricingMode string
DefaultDiscountFactor float64
PlatformRetryPolicy map[string]any
PlatformRateLimitPolicy map[string]any
PlatformPriority int
PlatformModelID string
BaseModelID string
CanonicalModelKey string
ProviderModelName string
ModelName string
ModelAlias string
ModelType string
DisplayName string
Capabilities map[string]any
CapabilityOverride map[string]any
BaseBillingConfig map[string]any
BillingConfig map[string]any
BillingConfigOverride map[string]any
PricingMode string
DiscountFactor float64
PlatformPricingRuleSetID string
ModelPricingRuleSetID string
ModelRetryPolicy map[string]any
ModelRateLimitPolicy map[string]any
ClientID string
QueueKey string
}
type RateLimitReservation struct {
ScopeType string
ScopeKey string
Metric string
Limit float64
Amount float64
WindowSeconds int
LeaseTTLSeconds int
}
type RateLimitResult struct {
LeaseIDs []string
}
type CreateTaskAttemptInput struct {
TaskID string
AttemptNo int
PlatformID string
PlatformModelID string
ClientID string
QueueKey string
Status string
Simulated bool
RequestSnapshot map[string]any
}
type FinishTaskAttemptInput struct {
AttemptID string
Status string
Retryable bool
ResponseSnapshot map[string]any
ErrorCode string
ErrorMessage string
}
apps/api/internal/store/tasks_runtime.go
+217
View File
@@ -0,0 +1,217 @@
package store
import (
"context"
"encoding/json"
)
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
_, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'running',
model_type = NULLIF($2, ''),
normalized_request = $3::jsonb,
locked_at = now(),
heartbeat_at = now(),
updated_at = now()
WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
return err
}
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
tx, err := s.pool.Begin(ctx)
if err != nil {
return "", err
}
defer tx.Rollback(ctx)
var attemptID string
err = tx.QueryRow(ctx, `
INSERT INTO gateway_task_attempts (
task_id, attempt_no, platform_id, platform_model_id, client_id, queue_key,
status, simulated, request_snapshot
)
VALUES (
$1::uuid, $2, NULLIF($3, '')::uuid, NULLIF($4, '')::uuid, NULLIF($5, ''), $6,
$7, $8, $9::jsonb
)
RETURNING id::text`,
input.TaskID,
input.AttemptNo,
input.PlatformID,
input.PlatformModelID,
input.ClientID,
input.QueueKey,
firstNonEmpty(input.Status, "running"),
input.Simulated,
string(requestJSON),
).Scan(&attemptID)
if err != nil {
return "", err
}
if _, err := tx.Exec(ctx, `
UPDATE gateway_tasks
SET attempt_count = GREATEST(attempt_count, $2), updated_at = now()
WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
return "", err
}
return attemptID, tx.Commit(ctx)
}
func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error {
responseJSON, _ := json.Marshal(emptyObjectIfNil(input.ResponseSnapshot))
_, err := s.pool.Exec(ctx, `
UPDATE gateway_task_attempts
SET status = $2,
retryable = $3,
response_snapshot = $4::jsonb,
error_code = NULLIF($5, ''),
error_message = NULLIF($6, ''),
finished_at = now()
WHERE id = $1::uuid`,
input.AttemptID,
input.Status,
input.Retryable,
string(responseJSON),
input.ErrorCode,
input.ErrorMessage,
)
return err
}
func (s *Store) FinishTaskSuccess(ctx context.Context, taskID string, result map[string]any, billings []any) (GatewayTask, error) {
resultJSON, _ := json.Marshal(emptyObjectIfNil(result))
billingsJSON, _ := json.Marshal(billings)
if _, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'succeeded',
result = $2::jsonb,
billings = $3::jsonb,
error = NULL,
error_code = NULL,
error_message = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`, taskID, string(resultJSON), string(billingsJSON)); err != nil {
return GatewayTask{}, err
}
return s.GetTask(ctx, taskID)
}
func (s *Store) FinishTaskFailure(ctx context.Context, taskID string, code string, message string) (GatewayTask, error) {
if _, err := s.pool.Exec(ctx, `
UPDATE gateway_tasks
SET status = 'failed',
error = NULLIF($2, ''),
error_code = NULLIF($3, ''),
error_message = NULLIF($2, ''),
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid`, taskID, message, code); err != nil {
return GatewayTask{}, err
}
return s.GetTask(ctx, taskID)
}
func (s *Store) AddTaskEvent(ctx context.Context, taskID string, eventType string, status string, phase string, progress float64, message string, payload map[string]any, simulated bool) (TaskEvent, error) {
payloadJSON, _ := json.Marshal(emptyObjectIfNil(payload))
var event TaskEvent
var payloadBytes []byte
err := s.pool.QueryRow(ctx, `
WITH next_seq AS (
SELECT COALESCE(MAX(seq), 0) + 1 AS seq
FROM gateway_task_events
WHERE task_id = $1::uuid
)
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
SELECT $1::uuid, next_seq.seq, $2, NULLIF($3, ''), NULLIF($4, ''), $5, NULLIF($6, ''), $7::jsonb, $8
FROM next_seq
RETURNING id::text, task_id::text, seq, event_type, COALESCE(status, ''), COALESCE(phase, ''),
COALESCE(progress, 0)::float8, COALESCE(message, ''), payload, simulated, created_at`,
taskID,
eventType,
status,
phase,
progress,
message,
string(payloadJSON),
simulated,
).Scan(
&event.ID,
&event.TaskID,
&event.Seq,
&event.EventType,
&event.Status,
&event.Phase,
&event.Progress,
&event.Message,
&payloadBytes,
&event.Simulated,
&event.CreatedAt,
)
if err != nil {
return TaskEvent{}, err
}
event.Payload = decodeObject(payloadBytes)
return event, nil
}
func (s *Store) QueueTaskCallback(ctx context.Context, event TaskEvent, callbackURL string) error {
if callbackURL == "" {
return nil
}
payloadJSON, _ := json.Marshal(map[string]any{
"taskId": event.TaskID,
"seq": event.Seq,
"eventType": event.EventType,
"status": event.Status,
"phase": event.Phase,
"progress": event.Progress,
"message": event.Message,
"payload": event.Payload,
"simulated": event.Simulated,
"createdAt": event.CreatedAt,
})
_, err := s.pool.Exec(ctx, `
INSERT INTO gateway_task_callback_outbox (task_id, event_id, seq, callback_url, payload)
VALUES ($1::uuid, $2::uuid, $3, $4, $5::jsonb)
ON CONFLICT (task_id, seq, callback_url) DO NOTHING`,
event.TaskID,
event.ID,
event.Seq,
callbackURL,
string(payloadJSON),
)
return err
}
func (s *Store) RecordClientAssignment(ctx context.Context, candidate RuntimeModelCandidate) error {
_, err := s.pool.Exec(ctx, `
INSERT INTO runtime_client_states (
client_id, platform_id, provider, method_name, queue_key, running_count, last_assigned_at
)
VALUES ($1, $2::uuid, $3, $4, $5, 1, now())
ON CONFLICT (client_id) DO UPDATE
SET running_count = runtime_client_states.running_count + 1,
last_assigned_at = now(),
updated_at = now()`,
candidate.ClientID,
candidate.PlatformID,
candidate.Provider,
candidate.ModelType,
candidate.QueueKey,
)
return err
}
func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastError string) error {
_, err := s.pool.Exec(ctx, `
UPDATE runtime_client_states
SET running_count = GREATEST(running_count - 1, 0),
last_error = NULLIF($2, ''),
updated_at = now()
WHERE client_id = $1`, clientID, lastError)
return err
}
apps/api/internal/store/user_group_policy.go
+41
View File
@@ -0,0 +1,41 @@
package store
import (
"context"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
)
type UserGroupPolicy struct {
ID string
GroupKey string
RateLimitPolicy map[string]any
BillingDiscountPolicy map[string]any
}
func (s *Store) ResolveUserGroupPolicy(ctx context.Context, user *auth.User) (UserGroupPolicy, error) {
userGroupID := ""
if user != nil {
userGroupID = user.UserGroupID
}
var item UserGroupPolicy
var rateLimit []byte
var billing []byte
err := s.pool.QueryRow(ctx, `
SELECT id::text, group_key, rate_limit_policy, billing_discount_policy
FROM gateway_user_groups
WHERE status = 'active'
AND (($1 <> '' AND id = NULLIF($1, '')::uuid) OR ($1 = '' AND group_key = 'default'))
ORDER BY CASE WHEN id::text = $1 THEN 0 ELSE 1 END, priority ASC
LIMIT 1`, userGroupID).Scan(&item.ID, &item.GroupKey, &rateLimit, &billing)
if err != nil {
if err == pgx.ErrNoRows {
return UserGroupPolicy{}, nil
}
return UserGroupPolicy{}, err
}
item.RateLimitPolicy = decodeObject(rateLimit)
item.BillingDiscountPolicy = decodeObject(billing)
return item, nil
}