chore: initial commit
This commit is contained in:
@@ -0,0 +1,229 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user