151 lines
4.3 KiB
Go
151 lines
4.3 KiB
Go
// Package route 维护逻辑模型/能力槽位到候选端点的路由注册表。
|
||
// 端点来源：环境 SHENSHI_ENDPOINTS（JSON 数组）；未配置时回落 stub（骨架形态）。
|
||
// 冷却粒度（shenshi-gateway.md §4）：401/403 只冷却凭证；429/5xx/超时只冷却该端点绑定。
|
||
package route
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"os"
|
||
"sort"
|
||
"time"
|
||
|
||
"hunjing/shenshi/internal/transform"
|
||
)
|
||
|
||
// Endpoint 一个可路由的出站端点。
|
||
type Endpoint struct {
|
||
Name string
|
||
Type string // openai_compatible / local / stub（契约词典 §14.2）
|
||
Adaptor transform.Adaptor
|
||
Priority int
|
||
// 凭证指纹（凭证级冷却键）；stub/openai 兼容端点各自填充
|
||
CredentialKey string
|
||
// 绑定级冷却键的模型名部分
|
||
UpstreamModel string
|
||
}
|
||
|
||
// EndpointSpec 环境配置形态。
|
||
type EndpointSpec struct {
|
||
Type string `json:"type"`
|
||
Name string `json:"name"`
|
||
BaseURL string `json:"base_url"`
|
||
APIKey string `json:"api_key"`
|
||
Model string `json:"model"`
|
||
Priority int `json:"priority"`
|
||
}
|
||
|
||
// Registry 槽位 → 候选端点列表（按 Priority 升序尝试）。
|
||
type Registry struct {
|
||
bySlot map[string][]Endpoint
|
||
// Cooldowns/FirstByteTimeoutMS 由 server 层共享使用
|
||
Cooldowns *Cooldowns
|
||
Sticky *Sticky
|
||
FirstByteTimeout time.Duration
|
||
}
|
||
|
||
// New 构建注册表：读 SHENSHI_ENDPOINTS（JSON），无配置时登记 stub。
|
||
func New() *Registry {
|
||
r := &Registry{
|
||
bySlot: make(map[string][]Endpoint),
|
||
Cooldowns: NewCooldowns(),
|
||
Sticky: NewSticky(),
|
||
FirstByteTimeout: 10 * time.Second,
|
||
}
|
||
specs := loadSpecs()
|
||
if len(specs) == 0 {
|
||
r.Register("chat", Endpoint{Type: "stub", Adaptor: transform.StubAdaptor{}, Priority: 100, UpstreamModel: "stub"})
|
||
return r
|
||
}
|
||
for _, spec := range specs {
|
||
if spec.Type != "openai_compatible" && spec.Type != "local" {
|
||
continue
|
||
}
|
||
adaptor := &transform.OpenAICompatAdaptor{
|
||
EndpointName: spec.Name,
|
||
BaseURL: spec.BaseURL,
|
||
APIKey: spec.APIKey,
|
||
UpstreamModel: spec.Model,
|
||
}
|
||
cred := spec.Name
|
||
if spec.APIKey != "" {
|
||
cred = spec.Name + ":" + fmt.Sprintf("%x", len(spec.APIKey)) // 不落密钥原文，仅指纹
|
||
}
|
||
r.Register("chat", Endpoint{
|
||
Name: spec.Name, Type: spec.Type, Adaptor: adaptor,
|
||
Priority: spec.Priority, CredentialKey: cred, UpstreamModel: spec.Model,
|
||
})
|
||
}
|
||
return r
|
||
}
|
||
|
||
func loadSpecs() []EndpointSpec {
|
||
raw := os.Getenv("SHENSHI_ENDPOINTS")
|
||
if raw == "" {
|
||
return nil
|
||
}
|
||
var specs []EndpointSpec
|
||
if err := json.Unmarshal([]byte(raw), &specs); err != nil {
|
||
return nil
|
||
}
|
||
return specs
|
||
}
|
||
|
||
// Reset 清空槽位（测试装配用）。
|
||
func (r *Registry) Reset(slot string) { delete(r.bySlot, slot) }
|
||
|
||
// Register 登记端点到指定能力槽位（契约词典 §14.1）。
|
||
func (r *Registry) Register(slot string, ep Endpoint) {
|
||
r.bySlot[slot] = append(r.bySlot[slot], ep)
|
||
}
|
||
|
||
// Candidates 返回槽位候选（Priority 升序；粘性端点置首；冷却中端点排后不剔除——
|
||
// 全部冷却时仍可尝试最优先者，避免死锁）。
|
||
func (r *Registry) Candidates(slot, runID string) ([]Endpoint, error) {
|
||
if slot == "" {
|
||
slot = "chat"
|
||
}
|
||
eps := append([]Endpoint(nil), r.bySlot[slot]...)
|
||
if len(eps) == 0 {
|
||
return nil, fmt.Errorf("槽位 %s 无可用端点", slot)
|
||
}
|
||
sort.SliceStable(eps, func(i, j int) bool { return eps[i].Priority < eps[j].Priority })
|
||
if runID != "" {
|
||
if stickyName := r.Sticky.Recall(runID); stickyName != "" {
|
||
for i, ep := range eps {
|
||
if ep.Name == stickyName {
|
||
eps = append([]Endpoint{ep}, append(eps[:i:i], eps[i+1:]...)...)
|
||
break
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// 冷却中的端点移到队尾（保留兜底能力）
|
||
cooling := 0
|
||
for i := 0; i < len(eps)-cooling; {
|
||
if r.Cooldowns.Active(bindingKey(eps[i])) || r.Cooldowns.Active(credKey(eps[i])) {
|
||
ep := eps[i]
|
||
eps = append(append(eps[:i:i], eps[i+1:]...), ep)
|
||
cooling++
|
||
continue
|
||
}
|
||
i++
|
||
}
|
||
return eps, nil
|
||
}
|
||
|
||
func bindingKey(ep Endpoint) string {
|
||
return "bind:" + ep.Name + ":" + ep.UpstreamModel
|
||
}
|
||
|
||
func credKey(ep Endpoint) string {
|
||
if ep.CredentialKey == "" {
|
||
return "cred:" + ep.Name
|
||
}
|
||
return "cred:" + ep.CredentialKey
|
||
}
|
||
|
||
// BindingCooldownKey / CredentialCooldownKey 供 server 层标记冷却。
|
||
func BindingCooldownKey(ep Endpoint) string { return bindingKey(ep) }
|
||
func CredentialCooldownKey(ep Endpoint) string { return credKey(ep) }
|