137 lines
3.7 KiB
Go
137 lines
3.7 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const (
|
|
discoveryTimeout = 30 * time.Second
|
|
discoveryMaxBody = 1 << 20
|
|
discoveryMaxPages = 20
|
|
)
|
|
|
|
func (s *Server) discoverModels(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
errJSON(w, 405, "method not allowed")
|
|
return
|
|
}
|
|
var c channel
|
|
if decode(w, r, &c) != nil {
|
|
errJSON(w, 400, "invalid channel")
|
|
return
|
|
}
|
|
u, e := url.Parse(c.BaseURL)
|
|
if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || strings.Contains(c.BaseURL, "#") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
|
|
errJSON(w, 400, "invalid channel")
|
|
return
|
|
}
|
|
if c.ID != "" {
|
|
var old string
|
|
e = s.db.QueryRowContext(r.Context(), "SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old)
|
|
if e == sql.ErrNoRows {
|
|
errJSON(w, 404, "not found")
|
|
return
|
|
}
|
|
if e != nil {
|
|
errJSON(w, 500, "database error")
|
|
return
|
|
}
|
|
if c.APIKey == "" {
|
|
c.APIKey = old
|
|
}
|
|
}
|
|
ctx, cancel := context.WithTimeout(r.Context(), discoveryTimeout)
|
|
defer cancel()
|
|
models, e := s.fetchModels(ctx, c)
|
|
if e != nil {
|
|
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
|
errJSON(w, 504, "model discovery timed out")
|
|
} else {
|
|
errJSON(w, 502, "model discovery failed")
|
|
}
|
|
return
|
|
}
|
|
writeJSON(w, map[string][]string{"models": models})
|
|
}
|
|
|
|
func (s *Server) fetchModels(ctx context.Context, c channel) ([]string, error) {
|
|
// Share the bounded transport, but never forward credentials across redirects.
|
|
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
|
endpoint := strings.TrimRight(c.BaseURL, "/") + "/models"
|
|
if c.Protocol == "anthropic" {
|
|
endpoint = strings.TrimRight(c.BaseURL, "/") + "/v1/models"
|
|
}
|
|
seen := map[string]bool{}
|
|
cursors := map[string]bool{}
|
|
after := ""
|
|
for page := 0; page < discoveryMaxPages; page++ {
|
|
u := endpoint
|
|
if after != "" {
|
|
u += "?" + url.Values{"after_id": {after}}.Encode()
|
|
}
|
|
req, e := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
req.Header.Set("Accept", "application/json")
|
|
if c.Protocol == "anthropic" {
|
|
req.Header.Set("x-api-key", c.APIKey)
|
|
req.Header.Set("anthropic-version", "2023-06-01")
|
|
} else {
|
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
|
}
|
|
res, e := client.Do(req)
|
|
if e != nil {
|
|
return nil, e
|
|
}
|
|
if res.StatusCode != http.StatusOK {
|
|
res.Body.Close()
|
|
return nil, errors.New("upstream status")
|
|
}
|
|
body, e := io.ReadAll(io.LimitReader(res.Body, discoveryMaxBody+1))
|
|
res.Body.Close()
|
|
if e != nil || len(body) > discoveryMaxBody {
|
|
return nil, errors.New("upstream body")
|
|
}
|
|
var result struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
HasMore bool `json:"has_more"`
|
|
LastID string `json:"last_id"`
|
|
}
|
|
if json.Unmarshal(body, &result) != nil || result.Data == nil {
|
|
return nil, errors.New("upstream models")
|
|
}
|
|
for _, model := range result.Data {
|
|
if strings.TrimSpace(model.ID) == "" {
|
|
return nil, errors.New("upstream model id")
|
|
}
|
|
seen[model.ID] = true
|
|
}
|
|
if c.Protocol != "anthropic" || !result.HasMore {
|
|
models := make([]string, 0, len(seen))
|
|
for id := range seen {
|
|
models = append(models, id)
|
|
}
|
|
sort.Strings(models)
|
|
return models, nil
|
|
}
|
|
if result.LastID == "" || cursors[result.LastID] {
|
|
return nil, errors.New("upstream pagination")
|
|
}
|
|
cursors[result.LastID] = true
|
|
after = result.LastID
|
|
}
|
|
return nil, errors.New("upstream page limit")
|
|
}
|