Files

250 lines
6.3 KiB
Go

package sub2api
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"strings"
)
func (c *Client) CreateAccount(ctx context.Context, req CreateAccountRequest) (AccountRef, error) {
var ref AccountRef
if err := c.postJSON(ctx, "/api/v1/admin/accounts", req, &ref); err != nil {
return AccountRef{}, err
}
return ref, nil
}
func (c *Client) BatchCreateAccounts(ctx context.Context, req BatchCreateAccountsRequest) ([]AccountRef, error) {
statusCode, _, body, err := c.perform(ctx, http.MethodPost, "/api/v1/admin/accounts/batch", req)
if err != nil {
return nil, err
}
if statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices {
return nil, newHTTPError(http.MethodPost, "/api/v1/admin/accounts/batch", statusCode, body)
}
models, err := decodeAccountRefs(body)
if err != nil {
return nil, fmt.Errorf("decode /api/v1/admin/accounts/batch response: %w", err)
}
return models, nil
}
func (c *Client) TestAccount(ctx context.Context, accountID, modelID string) (ProbeResult, error) {
path := "/api/v1/admin/accounts/" + accountID + "/test"
req := map[string]any{}
if strings.TrimSpace(modelID) != "" {
req["model_id"] = strings.TrimSpace(modelID)
}
statusCode, _, body, err := c.perform(ctx, http.MethodPost, path, req)
if err != nil {
return ProbeResult{}, err
}
if statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices {
return ProbeResult{}, newHTTPError(http.MethodPost, path, statusCode, body)
}
result, err := parseProbeResult(body)
if err != nil {
return ProbeResult{}, fmt.Errorf("parse %s sse: %w", path, err)
}
return result, nil
}
func (c *Client) GetAccountModels(ctx context.Context, accountID string) ([]AccountModel, error) {
path := "/api/v1/admin/accounts/" + accountID + "/models"
statusCode, _, body, err := c.perform(ctx, http.MethodGet, path, nil)
if err != nil {
return nil, err
}
if statusCode < http.StatusOK || statusCode >= http.StatusMultipleChoices {
return nil, newHTTPError(http.MethodGet, path, statusCode, body)
}
models, err := decodeAccountModels(body)
if err != nil {
return nil, fmt.Errorf("decode %s response: %w", path, err)
}
return models, nil
}
func decodeAccountRefs(body []byte) ([]AccountRef, error) {
var refs []AccountRef
if err := decodeEnvelopeObject(body, &refs); err == nil {
return refs, nil
}
var wrapper struct {
Data struct {
Items []AccountRef `json:"items"`
Results []struct {
ID json.RawMessage `json:"id"`
Name string `json:"name"`
Success bool `json:"success"`
} `json:"results"`
} `json:"data"`
}
if err := json.Unmarshal(body, &wrapper); err == nil {
if len(wrapper.Data.Items) > 0 {
return wrapper.Data.Items, nil
}
if len(wrapper.Data.Results) > 0 {
return decodeBatchAccountResults(wrapper.Data.Results)
}
}
var batch struct {
Results []struct {
ID json.RawMessage `json:"id"`
Name string `json:"name"`
Success bool `json:"success"`
} `json:"results"`
}
if err := json.Unmarshal(body, &batch); err != nil {
return nil, err
}
return decodeBatchAccountResults(batch.Results)
}
func decodeBatchAccountResults(results []struct {
ID json.RawMessage `json:"id"`
Name string `json:"name"`
Success bool `json:"success"`
}) ([]AccountRef, error) {
refs := make([]AccountRef, 0, len(results))
for _, item := range results {
if !item.Success {
continue
}
id, err := decodeFlexibleID(item.ID)
if err != nil {
return nil, err
}
refs = append(refs, AccountRef{ID: id, Name: item.Name})
}
return refs, nil
}
func decodeAccountModels(body []byte) ([]AccountModel, error) {
var models []AccountModel
if err := decodeEnvelopeObject(body, &models); err == nil {
return models, nil
}
var wrapper struct {
Data struct {
Items []AccountModel `json:"items"`
} `json:"data"`
}
if err := json.Unmarshal(body, &wrapper); err != nil {
return nil, err
}
return wrapper.Data.Items, nil
}
func parseProbeResult(body []byte) (ProbeResult, error) {
scanner := bufio.NewScanner(bytes.NewReader(body))
var payloads []string
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if strings.HasPrefix(line, "data:") {
payloads = append(payloads, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
}
}
if err := scanner.Err(); err != nil {
return ProbeResult{}, err
}
if len(payloads) == 0 {
return ProbeResult{}, fmt.Errorf("missing data event")
}
type probeEvent struct {
Type string `json:"type"`
Status string `json:"status"`
Message string `json:"message"`
Error string `json:"error"`
OK *bool `json:"ok"`
Success *bool `json:"success"`
}
var latest ProbeResult
sawProbeState := false
for _, payload := range payloads {
var event probeEvent
if err := json.Unmarshal([]byte(payload), &event); err != nil {
return ProbeResult{}, err
}
message := strings.TrimSpace(event.Message)
if message == "" {
message = strings.TrimSpace(event.Error)
}
eventType := strings.ToLower(strings.TrimSpace(event.Type))
if eventType == "error" || strings.TrimSpace(event.Error) != "" {
if message == "" {
message = "account probe returned an error event"
}
return ProbeResult{
OK: false,
Status: "failed",
Message: message,
}, nil
}
ok := false
switch {
case event.OK != nil:
ok = *event.OK
case event.Success != nil:
ok = *event.Success
default:
switch strings.ToLower(strings.TrimSpace(event.Status)) {
case "ok", "pass", "passed", "success", "succeeded":
ok = true
}
}
if eventType == "test_complete" {
return ProbeResult{
OK: ok,
Status: normalizeProbeStatus(event.Status, ok),
Message: message,
}, nil
}
if event.Status != "" || event.OK != nil || event.Success != nil {
latest = ProbeResult{
OK: ok,
Status: normalizeProbeStatus(event.Status, ok),
Message: message,
}
sawProbeState = true
}
}
if sawProbeState {
return latest, nil
}
return ProbeResult{}, fmt.Errorf("missing probe status event")
}
func normalizeProbeStatus(status string, ok bool) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "pass", "passed", "ok", "success", "succeeded":
return "passed"
case "fail", "failed", "error":
return "failed"
}
if ok {
return "passed"
}
return "failed"
}