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() }