initial upload
This commit is contained in:
Luxferre
+552
@@ -0,0 +1,552 @@
|
||||
// Dynagate LLM request proxying handlers
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var httpClient = &http.Client{}
|
||||
|
||||
func checkAuth(expectedToken string, r *http.Request) bool {
|
||||
if expectedToken == "" {
|
||||
return true
|
||||
}
|
||||
authHeader := r.Header.Get("Authorization")
|
||||
if !strings.HasPrefix(authHeader, "Bearer ") {
|
||||
return false
|
||||
}
|
||||
token := strings.TrimPrefix(authHeader, "Bearer ")
|
||||
return token == expectedToken
|
||||
}
|
||||
|
||||
func sendUnauthorized(w http.ResponseWriter) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Incorrect API key provided.",
|
||||
"type": "invalid_request_error",
|
||||
"param": nil,
|
||||
"code": "invalid_api_key",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func handleModels(cm *ConfigManager, expectedToken string) http.HandlerFunc {
|
||||
type ModelData struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
OwnedBy string `json:"owned_by"`
|
||||
}
|
||||
|
||||
type ModelsResponse struct {
|
||||
Object string `json:"object"`
|
||||
Data []ModelData `json:"data"`
|
||||
}
|
||||
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if !checkAuth(expectedToken, r) {
|
||||
sendUnauthorized(w)
|
||||
return
|
||||
}
|
||||
|
||||
if r.Method != http.MethodGet {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusMethodNotAllowed)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Method not allowed",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
models := cm.GetUniqueModels()
|
||||
data := make([]ModelData, len(models))
|
||||
for i, m := range models {
|
||||
data[i] = ModelData{
|
||||
ID: m,
|
||||
Object: "model",
|
||||
Created: 1686935002,
|
||||
OwnedBy: "dynagate",
|
||||
}
|
||||
}
|
||||
|
||||
resp := ModelsResponse{
|
||||
Object: "list",
|
||||
Data: data,
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(resp)
|
||||
}
|
||||
}
|
||||
|
||||
func handleChatCompletions(cm *ConfigManager, expectedToken string) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if !checkAuth(expectedToken, r) {
|
||||
sendUnauthorized(w)
|
||||
return
|
||||
}
|
||||
|
||||
if r.Method != http.MethodPost {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusMethodNotAllowed)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Method not allowed",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
bodyBytes, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Failed to read request body",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var bodyMap map[string]any
|
||||
if err := json.Unmarshal(bodyBytes, &bodyMap); err != nil {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Invalid JSON in request body",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var requestedModel string
|
||||
if m, ok := bodyMap["model"]; ok {
|
||||
if s, ok := m.(string); ok {
|
||||
requestedModel = s
|
||||
}
|
||||
}
|
||||
|
||||
configs := cm.GetConfigs()
|
||||
uniqueModels := cm.GetUniqueModels()
|
||||
|
||||
if len(configs) == 0 {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "No model configurations loaded",
|
||||
"type": "gateway_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
trialConfigs := getTrialConfigs(configs, uniqueModels, requestedModel)
|
||||
if len(trialConfigs) == 0 {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "No valid trial configuration candidates",
|
||||
"type": "gateway_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var isStream bool
|
||||
if s, ok := bodyMap["stream"]; ok {
|
||||
if b, ok := s.(bool); ok {
|
||||
isStream = b
|
||||
}
|
||||
}
|
||||
|
||||
log.Printf("Received completion request for model %q (stream=%t). Found %d config trials.", requestedModel, isStream, len(trialConfigs))
|
||||
|
||||
for i, trial := range trialConfigs {
|
||||
log.Printf("Trial %d/%d: model=%s endpoint=%s key_len=%d", i+1, len(trialConfigs), trial.Model, trial.Endpoint, len(trial.Key))
|
||||
|
||||
bodyMap["model"] = trial.Model
|
||||
modifiedBody, err := json.Marshal(bodyMap)
|
||||
if err != nil {
|
||||
log.Printf("Trial %d: Failed to marshal body for %s: %v", i+1, trial.Model, err)
|
||||
continue
|
||||
}
|
||||
|
||||
targetURL := buildURL(trial.Endpoint, r.URL.Path)
|
||||
|
||||
outReq, err := http.NewRequestWithContext(r.Context(), "POST", targetURL, bytes.NewReader(modifiedBody))
|
||||
if err != nil {
|
||||
log.Printf("Trial %d: Failed to create outgoing request to %s: %v", i+1, targetURL, err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Copy headers from incoming request, excluding host and auth
|
||||
for k, vv := range r.Header {
|
||||
kLower := strings.ToLower(k)
|
||||
if kLower == "authorization" || kLower == "host" || kLower == "content-length" {
|
||||
continue
|
||||
}
|
||||
for _, v := range vv {
|
||||
outReq.Header.Add(k, v)
|
||||
}
|
||||
}
|
||||
if trial.Key == "-blank-" {
|
||||
outReq.Header.Set("Authorization", "Bearer")
|
||||
} else if trial.Key != "" && trial.Key != "-" {
|
||||
outReq.Header.Set("Authorization", "Bearer "+trial.Key)
|
||||
}
|
||||
outReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
if trial.Extra != "" {
|
||||
var extraHeaders map[string]any
|
||||
if err := json.Unmarshal([]byte(trial.Extra), &extraHeaders); err != nil {
|
||||
log.Printf("Trial %d: Failed to parse extra headers JSON: %v", i+1, err)
|
||||
} else {
|
||||
for hk, hv := range extraHeaders {
|
||||
var valStr string
|
||||
switch v := hv.(type) {
|
||||
case string:
|
||||
valStr = v
|
||||
default:
|
||||
valStr = fmt.Sprintf("%v", v)
|
||||
}
|
||||
outReq.Header.Set(hk, valStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := httpClient.Do(outReq)
|
||||
if err != nil {
|
||||
log.Printf("Trial %d: Request to %s failed: %v", i+1, targetURL, err)
|
||||
continue
|
||||
}
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
|
||||
resp.Body.Close()
|
||||
log.Printf("Trial %d: Request to %s returned error status %d: %s", i+1, targetURL, resp.StatusCode, strings.TrimSpace(string(errBody)))
|
||||
continue
|
||||
}
|
||||
|
||||
log.Printf("Trial %d: Connection established with status %d. Proxying response.", i+1, resp.StatusCode)
|
||||
|
||||
if !isStream {
|
||||
respBody, readErr := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if readErr != nil {
|
||||
log.Printf("Trial %d: Failed to read non-streaming response body: %v", i+1, readErr)
|
||||
continue
|
||||
}
|
||||
|
||||
for k, vv := range resp.Header {
|
||||
for _, v := range vv {
|
||||
w.Header().Add(k, v)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(respBody)
|
||||
return
|
||||
}
|
||||
|
||||
// Handle streaming
|
||||
defer resp.Body.Close()
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
log.Printf("Trial %d: Flusher not supported on current ResponseWriter", i+1)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Response flusher not supported",
|
||||
"type": "gateway_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
for k, vv := range resp.Header {
|
||||
for _, v := range vv {
|
||||
w.Header().Add(k, v)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
flusher.Flush()
|
||||
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, rErr := resp.Body.Read(buf)
|
||||
if n > 0 {
|
||||
if _, wErr := w.Write(buf[:n]); wErr != nil {
|
||||
log.Printf("Trial %d: Client disconnected or write error: %v", i+1, wErr)
|
||||
return
|
||||
}
|
||||
flusher.Flush()
|
||||
}
|
||||
if rErr != nil {
|
||||
if rErr != io.EOF {
|
||||
log.Printf("Trial %d: Error reading response stream: %v", i+1, rErr)
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("All trials failed. Returning Bad Gateway.")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "All configured models and keys failed to respond.",
|
||||
"type": "gateway_error",
|
||||
"param": nil,
|
||||
"code": "all_endpoints_failed",
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func handleImageGenerations(cm *ConfigManager, expectedToken string) http.HandlerFunc {
|
||||
return func(w http.ResponseWriter, r *http.Request) {
|
||||
if !checkAuth(expectedToken, r) {
|
||||
sendUnauthorized(w)
|
||||
return
|
||||
}
|
||||
|
||||
if r.Method != http.MethodPost {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusMethodNotAllowed)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Method not allowed",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
bodyBytes, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Failed to read request body",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var bodyMap map[string]any
|
||||
if err := json.Unmarshal(bodyBytes, &bodyMap); err != nil {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "Invalid JSON in request body",
|
||||
"type": "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
var requestedModel string
|
||||
if m, ok := bodyMap["model"]; ok {
|
||||
if s, ok := m.(string); ok {
|
||||
requestedModel = s
|
||||
}
|
||||
}
|
||||
|
||||
configs := cm.GetConfigs()
|
||||
uniqueModels := cm.GetUniqueModels()
|
||||
|
||||
if len(configs) == 0 {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "No model configurations loaded",
|
||||
"type": "gateway_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
trialConfigs := getTrialConfigs(configs, uniqueModels, requestedModel)
|
||||
if len(trialConfigs) == 0 {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "No valid trial configuration candidates",
|
||||
"type": "gateway_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("Received image generation request for model %q. Found %d config trials.", requestedModel, len(trialConfigs))
|
||||
|
||||
for i, trial := range trialConfigs {
|
||||
log.Printf("Trial %d/%d: model=%s endpoint=%s key_len=%d", i+1, len(trialConfigs), trial.Model, trial.Endpoint, len(trial.Key))
|
||||
|
||||
bodyMap["model"] = trial.Model
|
||||
modifiedBody, err := json.Marshal(bodyMap)
|
||||
if err != nil {
|
||||
log.Printf("Trial %d: Failed to marshal body for %s: %v", i+1, trial.Model, err)
|
||||
continue
|
||||
}
|
||||
|
||||
targetURL := buildURL(trial.Endpoint, r.URL.Path)
|
||||
|
||||
outReq, err := http.NewRequestWithContext(r.Context(), "POST", targetURL, bytes.NewReader(modifiedBody))
|
||||
if err != nil {
|
||||
log.Printf("Trial %d: Failed to create outgoing request to %s: %v", i+1, targetURL, err)
|
||||
continue
|
||||
}
|
||||
|
||||
for k, vv := range r.Header {
|
||||
kLower := strings.ToLower(k)
|
||||
if kLower == "authorization" || kLower == "host" || kLower == "content-length" {
|
||||
continue
|
||||
}
|
||||
for _, v := range vv {
|
||||
outReq.Header.Add(k, v)
|
||||
}
|
||||
}
|
||||
if trial.Key == "-blank-" {
|
||||
outReq.Header.Set("Authorization", "Bearer")
|
||||
} else if trial.Key != "" && trial.Key != "-" {
|
||||
outReq.Header.Set("Authorization", "Bearer "+trial.Key)
|
||||
}
|
||||
outReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
if trial.Extra != "" {
|
||||
var extraHeaders map[string]any
|
||||
if err := json.Unmarshal([]byte(trial.Extra), &extraHeaders); err != nil {
|
||||
log.Printf("Trial %d: Failed to parse extra headers JSON: %v", i+1, err)
|
||||
} else {
|
||||
for hk, hv := range extraHeaders {
|
||||
var valStr string
|
||||
switch v := hv.(type) {
|
||||
case string:
|
||||
valStr = v
|
||||
default:
|
||||
valStr = fmt.Sprintf("%v", v)
|
||||
}
|
||||
outReq.Header.Set(hk, valStr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
resp, err := httpClient.Do(outReq)
|
||||
if err != nil {
|
||||
log.Printf("Trial %d: Request to %s failed: %v", i+1, targetURL, err)
|
||||
continue
|
||||
}
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
|
||||
resp.Body.Close()
|
||||
log.Printf("Trial %d: Request to %s returned error status %d: %s", i+1, targetURL, resp.StatusCode, strings.TrimSpace(string(errBody)))
|
||||
continue
|
||||
}
|
||||
|
||||
log.Printf("Trial %d: Connection established with status %d. Proxying response.", i+1, resp.StatusCode)
|
||||
|
||||
respBody, readErr := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
if readErr != nil {
|
||||
log.Printf("Trial %d: Failed to read response body: %v", i+1, readErr)
|
||||
continue
|
||||
}
|
||||
|
||||
for k, vv := range resp.Header {
|
||||
for _, v := range vv {
|
||||
w.Header().Add(k, v)
|
||||
}
|
||||
}
|
||||
w.WriteHeader(resp.StatusCode)
|
||||
_, _ = w.Write(respBody)
|
||||
return
|
||||
}
|
||||
|
||||
log.Printf("All trials failed. Returning Bad Gateway.")
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadGateway)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{
|
||||
"message": "All configured models and keys failed to respond.",
|
||||
"type": "gateway_error",
|
||||
"param": nil,
|
||||
"code": "all_endpoints_failed",
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
func getTrialConfigs(configs []ModelConfig, uniqueModels []string, requestedModel string) []ModelConfig {
|
||||
if len(configs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
reqIdx := -1
|
||||
for i, m := range uniqueModels {
|
||||
if m == requestedModel {
|
||||
reqIdx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
var orderedModels []string
|
||||
if reqIdx != -1 {
|
||||
orderedModels = append(orderedModels, uniqueModels[reqIdx:]...)
|
||||
orderedModels = append(orderedModels, uniqueModels[:reqIdx]...)
|
||||
} else {
|
||||
orderedModels = uniqueModels
|
||||
}
|
||||
|
||||
var trialConfigs []ModelConfig
|
||||
for _, modelName := range orderedModels {
|
||||
for _, cfg := range configs {
|
||||
if cfg.Model == modelName {
|
||||
trialConfigs = append(trialConfigs, cfg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return trialConfigs
|
||||
}
|
||||
|
||||
func buildURL(endpoint string, reqPath string) string {
|
||||
endpoint = strings.TrimSuffix(endpoint, "/")
|
||||
reqPath = strings.TrimPrefix(reqPath, "/")
|
||||
|
||||
if strings.HasSuffix(endpoint, "/v1") {
|
||||
if strings.HasPrefix(reqPath, "v1/") {
|
||||
reqPath = strings.TrimPrefix(reqPath, "v1/")
|
||||
}
|
||||
}
|
||||
|
||||
return endpoint + "/" + reqPath
|
||||
}
|
||||
Reference in New Issue
Block a user