239 lines
7.8 KiB
Go
239 lines
7.8 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"strings"
|
||
|
||
"github.com/agentmail/gateway/internal/db"
|
||
)
|
||
|
||
// ---------- 邮件场景下的可用模型 ----------
|
||
//
|
||
// 两张表，两种真相：
|
||
//
|
||
// agent_model_catalog —— 平台**上报**它当前看得见哪些模型（注册时整表替换）
|
||
// agent_allowed_models —— 管理员**选定**其中哪些能在邮件场景下用，rank 即优先级
|
||
//
|
||
// 为什么不合成一张带 allowed 标记的表：模型会从平台目录里消失（换了 provider 配置、
|
||
// 上游临时下线），那时整行被删掉就连带把管理员的选择也删了，模型回来还得重配一遍。
|
||
// 分开存之后，「选了什么」是持久的，目录只决定「这一项现在是否可用」。
|
||
//
|
||
// 为什么让平台上报而不是在 Gateway 里配一张静态表：模型清单是平台侧的事实 ——
|
||
// opencode 的 provider 配置、DSH 的 llm 适配器注册，都可能随时变。
|
||
// Gateway 猜不出来，猜错的后果是管理员在配置页选了一个平台其实调不到的模型。
|
||
|
||
// ModelRef 是一次「provider + model」路由。
|
||
type ModelRef struct {
|
||
Provider string `json:"provider"`
|
||
Model string `json:"model"`
|
||
}
|
||
|
||
// CatalogModel 是平台上报的一个可选模型。
|
||
type CatalogModel struct {
|
||
Provider string `json:"provider"`
|
||
Model string `json:"model"`
|
||
DisplayName string `json:"display_name,omitempty"`
|
||
// Allowed 表示它已被管理员选入邮件场景。
|
||
// 与目录合并后一起返回，前端才能画出「已勾选」的复选框。
|
||
Allowed bool `json:"allowed"`
|
||
// Rank 仅在 Allowed 为真时有意义，越小越先试。
|
||
Rank int `json:"rank,omitempty"`
|
||
}
|
||
|
||
// maxCatalogModels 限制单个 Agent 上报的模型数。
|
||
//
|
||
// 有平台会把上游的全部模型都列出来（实测 opencode 的一个 provider 就有几十个），
|
||
// 无上限的话一次注册能写进几千行，而配置页面上几千个复选框对人毫无用处。
|
||
const maxCatalogModels = 300
|
||
|
||
// ReplaceModelCatalog 整表替换某 Agent 上报的模型目录。
|
||
//
|
||
// 整表替换而非增量合并：目录是平台当前状态的快照，
|
||
// 增量合并会让已经下线的模型永远留在列表里，而那正是「选了却调不到」的来源。
|
||
//
|
||
// 事务包住删+插：中途失败留下一个空目录，会让配置页显示「该平台没有可用模型」
|
||
// 而管理员根本没做任何操作。
|
||
func ReplaceModelCatalog(ctx context.Context, agentName string, models []CatalogModel) error {
|
||
agentName = strings.TrimSpace(agentName)
|
||
if agentName == "" {
|
||
return nil
|
||
}
|
||
if len(models) > maxCatalogModels {
|
||
models = models[:maxCatalogModels]
|
||
}
|
||
|
||
tx, err := db.DB.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
if _, err := tx.ExecContext(ctx,
|
||
`DELETE FROM agent_model_catalog WHERE agent_name = $1`, agentName); err != nil {
|
||
return err
|
||
}
|
||
|
||
seen := map[string]bool{}
|
||
for _, m := range models {
|
||
p := strings.TrimSpace(m.Provider)
|
||
id := strings.TrimSpace(m.Model)
|
||
if p == "" || id == "" {
|
||
continue // 半条记录不如不要：它在配置页上是一个点不动的空复选框
|
||
}
|
||
key := p + "/" + id
|
||
if seen[key] {
|
||
continue
|
||
}
|
||
seen[key] = true
|
||
if _, err := tx.ExecContext(ctx,
|
||
`INSERT INTO agent_model_catalog (agent_name, provider, model, display_name, reported_at)
|
||
VALUES ($1, $2, $3, $4, NOW())`,
|
||
agentName, p, id, strings.TrimSpace(m.DisplayName)); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
// ListModelCatalog 返回某 Agent 的模型目录，并标出哪些已被选入邮件场景。
|
||
//
|
||
// LEFT JOIN 而不是两次查询：前端要的是一份「带勾选状态的清单」，
|
||
// 在 SQL 里合完比让前端自己对齐两个数组更难出错。
|
||
func ListModelCatalog(ctx context.Context, agentName string) ([]CatalogModel, error) {
|
||
rows, err := db.DB.QueryContext(ctx, `
|
||
SELECT c.provider, c.model, c.display_name,
|
||
CASE WHEN a.model IS NULL THEN 0 ELSE 1 END AS allowed,
|
||
COALESCE(a.rank, 0)
|
||
FROM agent_model_catalog c
|
||
LEFT JOIN agent_allowed_models a
|
||
ON a.agent_name = c.agent_name
|
||
AND a.provider = c.provider
|
||
AND a.model = c.model
|
||
WHERE c.agent_name = $1
|
||
ORDER BY c.provider, c.model
|
||
`, agentName)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
|
||
out := []CatalogModel{}
|
||
for rows.Next() {
|
||
var m CatalogModel
|
||
var allowed int
|
||
if err := rows.Scan(&m.Provider, &m.Model, &m.DisplayName, &allowed, &m.Rank); err != nil {
|
||
return nil, err
|
||
}
|
||
m.Allowed = allowed == 1
|
||
out = append(out, m)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// ListAllowedModels 按 rank 返回该 Agent 在邮件场景下可用的模型。
|
||
//
|
||
// **不与目录做 JOIN**：目录是平台上次注册时的快照，插件重启前可能已经过期。
|
||
// 真正能不能调通只有插件试过才知道 —— 这也正是插件要按顺序降级的原因。
|
||
// 在这里用目录过滤，只会把「目录暂时没上报但其实可用」的模型挡掉。
|
||
func ListAllowedModels(ctx context.Context, agentName string) ([]ModelRef, error) {
|
||
rows, err := db.DB.QueryContext(ctx, `
|
||
SELECT provider, model FROM agent_allowed_models
|
||
WHERE agent_name = $1
|
||
ORDER BY rank ASC, provider ASC, model ASC
|
||
`, agentName)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
|
||
out := []ModelRef{}
|
||
for rows.Next() {
|
||
var m ModelRef
|
||
if err := rows.Scan(&m.Provider, &m.Model); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, m)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// ListStaleAllowedModels 返回已选但已不在平台目录里的模型。
|
||
//
|
||
// 平台可能临时下线了某个模型（换了 provider 配置、上游故障），
|
||
// 而管理员的选择是持久的。界面上不显示这些项的话，管理员会以为自己
|
||
// 没选过它们 —— 而它们其实还在被插件尝试（ListAllowedModels 不与目录 JOIN）。
|
||
func ListStaleAllowedModels(ctx context.Context, agentName string) ([]ModelRef, error) {
|
||
rows, err := db.DB.QueryContext(ctx, `
|
||
SELECT a.provider, a.model
|
||
FROM agent_allowed_models a
|
||
WHERE a.agent_name = $1
|
||
AND NOT EXISTS (
|
||
SELECT 1 FROM agent_model_catalog c
|
||
WHERE c.agent_name = a.agent_name
|
||
AND c.provider = a.provider
|
||
AND c.model = a.model
|
||
)
|
||
ORDER BY a.rank ASC
|
||
`, agentName)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
defer rows.Close()
|
||
|
||
out := []ModelRef{}
|
||
for rows.Next() {
|
||
var m ModelRef
|
||
if err := rows.Scan(&m.Provider, &m.Model); err != nil {
|
||
return nil, err
|
||
}
|
||
out = append(out, m)
|
||
}
|
||
return out, rows.Err()
|
||
}
|
||
|
||
// SetAllowedModels 整表替换某 Agent 的邮件场景可用模型，入参顺序即优先级。
|
||
//
|
||
// 允许传空列表：那表示「不限定」——插件回退到平台自己的默认模型。
|
||
// 这与「一个都不许用」不同，后者等于让 Agent 彻底哑掉，不该是一次误删的后果。
|
||
func SetAllowedModels(ctx context.Context, agentName string, picks []ModelRef) error {
|
||
agentName = strings.TrimSpace(agentName)
|
||
if agentName == "" {
|
||
return nil
|
||
}
|
||
|
||
tx, err := db.DB.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
if _, err := tx.ExecContext(ctx,
|
||
`DELETE FROM agent_allowed_models WHERE agent_name = $1`, agentName); err != nil {
|
||
return err
|
||
}
|
||
|
||
rank := 0
|
||
seen := map[string]bool{}
|
||
for _, m := range picks {
|
||
p := strings.TrimSpace(m.Provider)
|
||
id := strings.TrimSpace(m.Model)
|
||
if p == "" || id == "" {
|
||
continue
|
||
}
|
||
key := p + "/" + id
|
||
if seen[key] {
|
||
// 重复项直接跳过而不是报错：它对最终顺序没有影响，
|
||
// 为一次无害的重复让整次保存失败只会让人以为配置没生效。
|
||
continue
|
||
}
|
||
seen[key] = true
|
||
if _, err := tx.ExecContext(ctx,
|
||
`INSERT INTO agent_allowed_models (agent_name, provider, model, rank)
|
||
VALUES ($1, $2, $3, $4)`, agentName, p, id, rank); err != nil {
|
||
return err
|
||
}
|
||
rank++
|
||
}
|
||
return tx.Commit()
|
||
}
|