63 lines
1.5 KiB
Go
63 lines
1.5 KiB
Go
|
package sub2api
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
)
|
||
|
|
|
||
|
|
type GatewayAccessCheckRequest struct {
|
||
|
|
APIKey string
|
||
|
|
ExpectedModel string
|
||
|
|
}
|
||
|
|
|
||
|
|
type GatewayAccessResult struct {
|
||
|
|
OK bool `json:"ok"`
|
||
|
|
StatusCode int `json:"status_code"`
|
||
|
|
Models []string `json:"models"`
|
||
|
|
HasExpectedModel bool `json:"has_expected_model"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *Client) CheckGatewayAccess(ctx context.Context, req GatewayAccessCheckRequest) (GatewayAccessResult, error) {
|
||
|
|
gatewayClient := *c
|
||
|
|
gatewayClient.apiKey = strings.TrimSpace(req.APIKey)
|
||
|
|
gatewayClient.bearerToken = ""
|
||
|
|
|
||
|
|
statusCode, _, body, err := gatewayClient.perform(ctx, http.MethodGet, "/v1/models", nil)
|
||
|
|
if err != nil {
|
||
|
|
return GatewayAccessResult{}, err
|
||
|
|
}
|
||
|
|
|
||
|
|
result := GatewayAccessResult{StatusCode: statusCode, OK: statusCode >= http.StatusOK && statusCode < http.StatusMultipleChoices}
|
||
|
|
if !result.OK {
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
result.Models = decodeGatewayModelIDs(body)
|
||
|
|
for _, modelID := range result.Models {
|
||
|
|
if modelID == strings.TrimSpace(req.ExpectedModel) {
|
||
|
|
result.HasExpectedModel = true
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func decodeGatewayModelIDs(body []byte) []string {
|
||
|
|
var payload struct {
|
||
|
|
Data []struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
} `json:"data"`
|
||
|
|
}
|
||
|
|
if err := json.Unmarshal(body, &payload); err == nil && len(payload.Data) > 0 {
|
||
|
|
models := make([]string, 0, len(payload.Data))
|
||
|
|
for _, item := range payload.Data {
|
||
|
|
if id := strings.TrimSpace(item.ID); id != "" {
|
||
|
|
models = append(models, id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return models
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|