230 lines
6.9 KiB
Go
230 lines
6.9 KiB
Go
package server
|
|
|
|
import (
|
|
"ai-router/internal/protocol"
|
|
"bytes"
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
const modelTestTimeout = 30 * time.Second
|
|
const modelTestMaxBody = 1 << 20
|
|
|
|
type modelTestResult struct {
|
|
OK bool `json:"ok"`
|
|
DurationMS int64 `json:"duration_ms"`
|
|
Status int `json:"status"`
|
|
Error string `json:"error,omitempty"`
|
|
}
|
|
|
|
// testModel is an administrative diagnostic, excluded from call logs and usage
|
|
// statistics. It creates no API key; upstream providers may still bill the probe.
|
|
func (s *Server) testModel(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
errJSON(w, 405, "method not allowed")
|
|
return
|
|
}
|
|
var input struct {
|
|
Model string `json:"model"`
|
|
Prompt *string `json:"prompt"`
|
|
}
|
|
if decode(w, r, &input) != nil {
|
|
errJSON(w, 400, "invalid request")
|
|
return
|
|
}
|
|
name, alias, ok := strings.Cut(input.Model, "/")
|
|
prompt := "Reply with OK."
|
|
if input.Prompt != nil {
|
|
prompt = *input.Prompt
|
|
}
|
|
if !ok || name == "" || alias == "" || strings.TrimSpace(prompt) == "" || utf8.RuneCountInString(prompt) > 2000 {
|
|
errJSON(w, 400, "invalid model or prompt")
|
|
return
|
|
}
|
|
ctx, cancel := context.WithTimeout(r.Context(), modelTestTimeout)
|
|
defer cancel()
|
|
var c channel
|
|
var models string
|
|
e := s.db.QueryRowContext(ctx, "SELECT base_url,protocol,enabled,models,api_key FROM channels WHERE name=?", name).Scan(&c.BaseURL, &c.Protocol, &c.Enabled, &models, &c.APIKey)
|
|
if e == sql.ErrNoRows {
|
|
errJSON(w, 403, "model or channel unavailable")
|
|
return
|
|
}
|
|
if e != nil || json.Unmarshal([]byte(models), &c.Models) != nil {
|
|
errJSON(w, 500, "database error")
|
|
return
|
|
}
|
|
actual, mapped := c.Models[alias]
|
|
if !c.Enabled || !mapped || actual == "" {
|
|
errJSON(w, 403, "model or channel unavailable")
|
|
return
|
|
}
|
|
writeJSON(w, s.probeModel(ctx, c, actual, prompt))
|
|
}
|
|
|
|
func (s *Server) probeModel(ctx context.Context, c channel, actual, prompt string) (result modelTestResult) {
|
|
start := time.Now()
|
|
defer func() {
|
|
result.DurationMS = time.Since(start).Milliseconds()
|
|
if ctx.Err() != nil {
|
|
result.OK = false
|
|
result.Error = "model test canceled"
|
|
if ctx.Err() == context.DeadlineExceeded {
|
|
result.Error = "model test timed out"
|
|
}
|
|
}
|
|
}()
|
|
result.Error = "upstream request failed"
|
|
up := protocol.Kind(c.Protocol)
|
|
body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"role": "user", "content": prompt}}, "stream": false, "max_completion_tokens": 128})
|
|
payload, e := protocol.ConvertRequest(protocol.Chat, up, body, actual)
|
|
if e != nil {
|
|
result.Error = "invalid upstream configuration"
|
|
return
|
|
}
|
|
suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
|
|
req, e := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+suffix, bytes.NewReader(payload))
|
|
if e != nil {
|
|
result.Error = "invalid upstream configuration"
|
|
return
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
if up == 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)
|
|
}
|
|
client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
|
res, e := client.Do(req)
|
|
if e != nil {
|
|
if errors.Is(e, context.DeadlineExceeded) {
|
|
result.Error = "model test timed out"
|
|
}
|
|
return
|
|
}
|
|
defer res.Body.Close()
|
|
result.Status = res.StatusCode
|
|
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
|
result.Error = "upstream rejected request"
|
|
return
|
|
}
|
|
body, e = io.ReadAll(io.LimitReader(res.Body, modelTestMaxBody+1))
|
|
if len(body) > modelTestMaxBody {
|
|
result.Error = "upstream response too large"
|
|
return
|
|
}
|
|
if e != nil || !validModelTestResponse(up, body) {
|
|
result.Error = "invalid upstream response"
|
|
return
|
|
}
|
|
result.OK, result.Error = true, ""
|
|
return
|
|
}
|
|
|
|
func validModelTestResponse(kind protocol.Kind, body []byte) bool {
|
|
var r struct {
|
|
ID string `json:"id"`
|
|
Object string `json:"object"`
|
|
Type string `json:"type"`
|
|
Role string `json:"role"`
|
|
Error json.RawMessage `json:"error"`
|
|
Status string `json:"status"`
|
|
StopReason string `json:"stop_reason"`
|
|
Incomplete struct {
|
|
Reason string `json:"reason"`
|
|
} `json:"incomplete_details"`
|
|
Output []json.RawMessage `json:"output"`
|
|
Content []json.RawMessage `json:"content"`
|
|
Choices []struct {
|
|
Finish string `json:"finish_reason"`
|
|
Message struct {
|
|
Role string `json:"role"`
|
|
Content *string `json:"content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
if json.Unmarshal(body, &r) != nil || r.ID == "" || (len(r.Error) != 0 && string(r.Error) != "null") {
|
|
return false
|
|
}
|
|
switch kind {
|
|
case protocol.Chat:
|
|
if r.Object != "chat.completion" || len(r.Choices) == 0 {
|
|
return false
|
|
}
|
|
for _, c := range r.Choices {
|
|
if c.Message.Role != "assistant" || (c.Finish != "stop" && c.Finish != "length") || (c.Message.Content == nil && c.Finish != "length") {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
case protocol.Responses:
|
|
if r.Object != "response" || r.Output == nil {
|
|
return false
|
|
}
|
|
exhausted := r.Status == "incomplete" && r.Incomplete.Reason == "max_output_tokens"
|
|
if !exhausted && (r.Status != "completed" || len(r.Output) == 0) {
|
|
return false
|
|
}
|
|
for _, raw := range r.Output {
|
|
var item struct {
|
|
Type string `json:"type"`
|
|
Role string `json:"role"`
|
|
Content []json.RawMessage `json:"content"`
|
|
Summary []json.RawMessage `json:"summary"`
|
|
}
|
|
if json.Unmarshal(raw, &item) != nil {
|
|
return false
|
|
}
|
|
if item.Type == "reasoning" && item.Summary != nil {
|
|
continue
|
|
}
|
|
if item.Type != "message" || item.Role != "assistant" || !validProbeBlocks(item.Content, "output_text") {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
case protocol.Anthropic:
|
|
if r.Type != "message" || r.Role != "assistant" || r.Content == nil {
|
|
return false
|
|
}
|
|
if r.StopReason == "max_tokens" && len(r.Content) == 0 {
|
|
return true
|
|
}
|
|
return (r.StopReason == "end_turn" || r.StopReason == "max_tokens") && validProbeBlocks(r.Content, "text")
|
|
}
|
|
return false
|
|
}
|
|
|
|
func validProbeBlocks(blocks []json.RawMessage, typ string) bool {
|
|
if len(blocks) == 0 {
|
|
return false
|
|
}
|
|
for _, raw := range blocks {
|
|
var b struct {
|
|
Type string `json:"type"`
|
|
Text *string `json:"text"`
|
|
Thinking *string `json:"thinking"`
|
|
Signature *string `json:"signature"`
|
|
Data *string `json:"data"`
|
|
}
|
|
if json.Unmarshal(raw, &b) != nil {
|
|
return false
|
|
}
|
|
if typ == "text" && ((b.Type == "thinking" && b.Thinking != nil && b.Signature != nil) || (b.Type == "redacted_thinking" && b.Data != nil)) {
|
|
continue
|
|
}
|
|
if b.Type != typ || b.Text == nil {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|