feat(commands): add Telegram model mode controls
This commit is contained in:
parent
bd56e10bb8
commit
33cf828760
6 changed files with 643 additions and 20 deletions
|
|
@ -58,6 +58,8 @@ type AgentLoop struct {
|
|||
hookRuntime hookRuntime
|
||||
steering *steeringQueue
|
||||
pendingSkills sync.Map
|
||||
pendingModelOverrides sync.Map
|
||||
sessionModelOverrides sync.Map
|
||||
mu sync.RWMutex
|
||||
|
||||
// Concurrent turn management (from HEAD)
|
||||
|
|
@ -82,6 +84,7 @@ type processOptions struct {
|
|||
SenderDisplayName string // Current sender display name for dynamic context
|
||||
UserMessage string // User message content (may include prefix)
|
||||
ForcedSkills []string // Skills explicitly requested for this message
|
||||
ForcedModel string // Model explicitly forced for this message/session
|
||||
SystemPromptOverride string // Override the default system prompt (Used by SubTurns)
|
||||
Media []string // media:// refs from inbound message
|
||||
InitialSteeringMessages []providers.Message // Steering messages from refactor/agent
|
||||
|
|
@ -1418,6 +1421,22 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
|||
return response, nil
|
||||
}
|
||||
|
||||
if pendingModel := al.takePendingModelOverride(opts.SessionKey); pendingModel != "" {
|
||||
opts.ForcedModel = pendingModel
|
||||
logger.InfoCF("agent", "Applying pending model override",
|
||||
map[string]any{
|
||||
"session_key": opts.SessionKey,
|
||||
"model": pendingModel,
|
||||
})
|
||||
} else if sessionModel := al.getSessionModelOverride(opts.SessionKey); sessionModel != "" {
|
||||
opts.ForcedModel = sessionModel
|
||||
logger.DebugCF("agent", "Applying session model override",
|
||||
map[string]any{
|
||||
"session_key": opts.SessionKey,
|
||||
"model": sessionModel,
|
||||
})
|
||||
}
|
||||
|
||||
if pending := al.takePendingSkills(opts.SessionKey); len(pending) > 0 {
|
||||
opts.ForcedSkills = append(opts.ForcedSkills, pending...)
|
||||
logger.InfoCF("agent", "Applying pending skill override",
|
||||
|
|
@ -1784,11 +1803,35 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
|
|||
ts.ingestMessage(turnCtx, al, rootMsg)
|
||||
}
|
||||
|
||||
activeCandidates, activeModel, usedLight := al.selectCandidates(ts.agent, ts.userMessage, messages)
|
||||
activeCandidates := ts.agent.Candidates
|
||||
activeModel := resolvedCandidateModel(ts.agent.Candidates, ts.agent.Model)
|
||||
activeProvider := ts.agent.Provider
|
||||
activeThinkingLevel := ts.agent.ThinkingLevel
|
||||
var usedLight bool
|
||||
if forcedModel := strings.TrimSpace(ts.opts.ForcedModel); forcedModel != "" {
|
||||
forcedCfg, forcedProvider, forcedCandidates, err := al.resolveModelSelection(cfg, ts.agent, forcedModel)
|
||||
if err != nil {
|
||||
turnStatus = TurnEndStatusError
|
||||
return turnResult{}, err
|
||||
}
|
||||
activeCandidates = forcedCandidates
|
||||
activeModel = resolvedCandidateModel(forcedCandidates, forcedModel)
|
||||
activeProvider = forcedProvider
|
||||
activeThinkingLevel = parseThinkingLevel(forcedCfg.ThinkingLevel)
|
||||
if stateful, ok := activeProvider.(providers.StatefulProvider); ok {
|
||||
defer stateful.Close()
|
||||
}
|
||||
logger.InfoCF("agent", "Forced model override selected",
|
||||
map[string]any{
|
||||
"agent_id": ts.agent.ID,
|
||||
"model": activeModel,
|
||||
})
|
||||
} else {
|
||||
activeCandidates, activeModel, usedLight = al.selectCandidates(ts.agent, ts.userMessage, messages)
|
||||
if usedLight && ts.agent.LightProvider != nil {
|
||||
activeProvider = ts.agent.LightProvider
|
||||
}
|
||||
}
|
||||
pendingMessages := append([]providers.Message(nil), ts.opts.InitialSteeringMessages...)
|
||||
var finalContent string
|
||||
|
||||
|
|
@ -1893,7 +1936,7 @@ turnLoop:
|
|||
hasWebSearch &&
|
||||
func() bool {
|
||||
// Check if provider supports native search
|
||||
if ns, ok := ts.agent.Provider.(interface{ SupportsNativeSearch() bool }); ok {
|
||||
if ns, ok := activeProvider.(interface{ SupportsNativeSearch() bool }); ok {
|
||||
return ns.SupportsNativeSearch()
|
||||
}
|
||||
return false
|
||||
|
|
@ -1933,12 +1976,12 @@ turnLoop:
|
|||
if useNativeSearch {
|
||||
llmOpts["native_search"] = true
|
||||
}
|
||||
if ts.agent.ThinkingLevel != ThinkingOff {
|
||||
if tc, ok := ts.agent.Provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() {
|
||||
llmOpts["thinking_level"] = string(ts.agent.ThinkingLevel)
|
||||
if activeThinkingLevel != ThinkingOff {
|
||||
if tc, ok := activeProvider.(providers.ThinkingCapable); ok && tc.SupportsThinking() {
|
||||
llmOpts["thinking_level"] = string(activeThinkingLevel)
|
||||
} else {
|
||||
logger.WarnCF("agent", "thinking_level is set but current provider does not support it, ignoring",
|
||||
map[string]any{"agent_id": ts.agent.ID, "thinking_level": string(ts.agent.ThinkingLevel)})
|
||||
map[string]any{"agent_id": ts.agent.ID, "thinking_level": string(activeThinkingLevel)})
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -3187,6 +3230,35 @@ func (al *AgentLoop) applyExplicitSkillCommand(
|
|||
return true, false, ""
|
||||
}
|
||||
|
||||
func (al *AgentLoop) resolveModelSelection(cfg *config.Config, agent *AgentInstance, modelName string) (*config.ModelConfig, providers.LLMProvider, []providers.FallbackCandidate, error) {
|
||||
if cfg == nil {
|
||||
return nil, nil, nil, fmt.Errorf("config is nil")
|
||||
}
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if modelName == "" {
|
||||
return nil, nil, nil, fmt.Errorf("model name is required")
|
||||
}
|
||||
workspace := ""
|
||||
fallbacks := []string(nil)
|
||||
if agent != nil {
|
||||
workspace = agent.Workspace
|
||||
fallbacks = agent.Fallbacks
|
||||
}
|
||||
modelCfg, err := resolvedModelConfig(cfg, modelName, workspace)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
provider, _, err := providers.CreateProviderFromConfig(modelCfg)
|
||||
if err != nil {
|
||||
return nil, nil, nil, fmt.Errorf("failed to initialize model %q: %w", modelName, err)
|
||||
}
|
||||
candidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, modelName, fallbacks)
|
||||
if len(candidates) == 0 {
|
||||
return nil, nil, nil, fmt.Errorf("model %q did not resolve to any provider candidates", modelName)
|
||||
}
|
||||
return modelCfg, provider, candidates, nil
|
||||
}
|
||||
|
||||
func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOptions) *commands.Runtime {
|
||||
registry := al.GetRegistry()
|
||||
cfg := al.GetConfig()
|
||||
|
|
@ -3265,6 +3337,42 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
|||
return oldModel, nil
|
||||
}
|
||||
|
||||
validateSessionModel := func(value string) error {
|
||||
if agent == nil {
|
||||
return fmt.Errorf("session model overrides unavailable in current context")
|
||||
}
|
||||
_, _, _, err := al.resolveModelSelection(cfg, agent, value)
|
||||
return err
|
||||
}
|
||||
|
||||
rt.GetSessionModelMode = func() (string, string) {
|
||||
if opts == nil || strings.TrimSpace(opts.SessionKey) == "" {
|
||||
return "", ""
|
||||
}
|
||||
return al.getSessionModelOverride(opts.SessionKey), al.getPendingModelOverride(opts.SessionKey)
|
||||
}
|
||||
rt.SetSessionModelMode = func(value string) error {
|
||||
if opts == nil || strings.TrimSpace(opts.SessionKey) == "" {
|
||||
return fmt.Errorf("session model overrides unavailable in current context")
|
||||
}
|
||||
if err := validateSessionModel(value); err != nil {
|
||||
return err
|
||||
}
|
||||
al.setSessionModelOverride(opts.SessionKey, value)
|
||||
al.clearPendingModelOverride(opts.SessionKey)
|
||||
return nil
|
||||
}
|
||||
rt.ArmNextModelMode = func(value string) error {
|
||||
if opts == nil || strings.TrimSpace(opts.SessionKey) == "" {
|
||||
return fmt.Errorf("session model overrides unavailable in current context")
|
||||
}
|
||||
if err := validateSessionModel(value); err != nil {
|
||||
return err
|
||||
}
|
||||
al.setPendingModelOverride(opts.SessionKey, value)
|
||||
return nil
|
||||
}
|
||||
|
||||
rt.ClearHistory = func() error {
|
||||
if opts == nil {
|
||||
return fmt.Errorf("process options not available")
|
||||
|
|
@ -3302,6 +3410,96 @@ func buildUseCommandHelp(agent *AgentInstance) string {
|
|||
)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) setSessionModelOverride(sessionKey, modelName string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if sessionKey == "" {
|
||||
return
|
||||
}
|
||||
if modelName == "" {
|
||||
al.sessionModelOverrides.Delete(sessionKey)
|
||||
return
|
||||
}
|
||||
al.sessionModelOverrides.Store(sessionKey, modelName)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) getSessionModelOverride(sessionKey string) string {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return ""
|
||||
}
|
||||
value, ok := al.sessionModelOverrides.Load(sessionKey)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
modelName, ok := value.(string)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(modelName)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) clearSessionModelOverride(sessionKey string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return
|
||||
}
|
||||
al.sessionModelOverrides.Delete(sessionKey)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) setPendingModelOverride(sessionKey, modelName string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if sessionKey == "" {
|
||||
return
|
||||
}
|
||||
if modelName == "" {
|
||||
al.pendingModelOverrides.Delete(sessionKey)
|
||||
return
|
||||
}
|
||||
al.pendingModelOverrides.Store(sessionKey, modelName)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) getPendingModelOverride(sessionKey string) string {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return ""
|
||||
}
|
||||
value, ok := al.pendingModelOverrides.Load(sessionKey)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
modelName, ok := value.(string)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(modelName)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) takePendingModelOverride(sessionKey string) string {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return ""
|
||||
}
|
||||
value, ok := al.pendingModelOverrides.LoadAndDelete(sessionKey)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
modelName, ok := value.(string)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
return strings.TrimSpace(modelName)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) clearPendingModelOverride(sessionKey string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" {
|
||||
return
|
||||
}
|
||||
al.pendingModelOverrides.Delete(sessionKey)
|
||||
}
|
||||
|
||||
func (al *AgentLoop) setPendingSkills(sessionKey string, skillNames []string) {
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if sessionKey == "" || len(skillNames) == 0 {
|
||||
|
|
|
|||
|
|
@ -1753,6 +1753,126 @@ func TestProcessMessage_SwitchModelRoutesSubsequentRequestsToSelectedProvider(t
|
|||
}
|
||||
}
|
||||
|
||||
func TestProcessMessage_BoostForcesPaidModelForOneTurn(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
heavyCalls := 0
|
||||
heavyServer := newStrictChatCompletionTestServer(
|
||||
t,
|
||||
"heavy",
|
||||
"gemini-2.5-flash",
|
||||
"heavy reply",
|
||||
&heavyCalls,
|
||||
)
|
||||
defer heavyServer.Close()
|
||||
|
||||
lightCalls := 0
|
||||
lightServer := newStrictChatCompletionTestServer(
|
||||
t,
|
||||
"light",
|
||||
"qwen2.5:0.5b",
|
||||
"light reply",
|
||||
&lightCalls,
|
||||
)
|
||||
defer lightServer.Close()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
ModelName: "gemini-main",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
Routing: &config.RoutingConfig{
|
||||
Enabled: true,
|
||||
LightModel: "qwen-light",
|
||||
Threshold: 0.99,
|
||||
},
|
||||
},
|
||||
},
|
||||
ModelList: []*config.ModelConfig{
|
||||
{
|
||||
ModelName: "gemini-main",
|
||||
Model: "gemini/gemini-2.5-flash",
|
||||
APIBase: heavyServer.URL,
|
||||
APIKeys: config.SimpleSecureStrings("heavy-key"),
|
||||
},
|
||||
{
|
||||
ModelName: "qwen-light",
|
||||
Model: "ollama/qwen2.5:0.5b",
|
||||
APIBase: lightServer.URL,
|
||||
APIKeys: config.SimpleSecureStrings("light-key"),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
provider, _, err := providers.CreateProvider(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateProvider() error = %v", err)
|
||||
}
|
||||
al := NewAgentLoop(cfg, msgBus, provider)
|
||||
helper := testHelper{al: al}
|
||||
|
||||
boostResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{
|
||||
Channel: "telegram",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "/boost",
|
||||
Peer: bus.Peer{
|
||||
Kind: "direct",
|
||||
ID: "user1",
|
||||
},
|
||||
})
|
||||
if !strings.Contains(boostResp, "Boost armed. Next message will use gemini-main.") {
|
||||
t.Fatalf("unexpected /boost reply: %q", boostResp)
|
||||
}
|
||||
|
||||
firstResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{
|
||||
Channel: "telegram",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "hello after boost",
|
||||
Peer: bus.Peer{
|
||||
Kind: "direct",
|
||||
ID: "user1",
|
||||
},
|
||||
})
|
||||
if firstResp != "heavy reply" {
|
||||
t.Fatalf("unexpected boosted response: %q", firstResp)
|
||||
}
|
||||
if heavyCalls != 1 {
|
||||
t.Fatalf("heavy calls after boost = %d, want 1", heavyCalls)
|
||||
}
|
||||
if lightCalls != 0 {
|
||||
t.Fatalf("light calls after boost = %d, want 0", lightCalls)
|
||||
}
|
||||
|
||||
secondResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{
|
||||
Channel: "telegram",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "hello after boost consumed",
|
||||
Peer: bus.Peer{
|
||||
Kind: "direct",
|
||||
ID: "user1",
|
||||
},
|
||||
})
|
||||
if secondResp != "light reply" {
|
||||
t.Fatalf("unexpected post-boost response: %q", secondResp)
|
||||
}
|
||||
if heavyCalls != 1 {
|
||||
t.Fatalf("heavy calls after second message = %d, want 1", heavyCalls)
|
||||
}
|
||||
if lightCalls != 1 {
|
||||
t.Fatalf("light calls after second message = %d, want 1", lightCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessMessage_ModelRoutingUsesLightProvider(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -12,6 +12,10 @@ func BuiltinDefinitions() []Definition {
|
|||
listCommand(),
|
||||
useCommand(),
|
||||
switchCommand(),
|
||||
boostCommand(),
|
||||
paidCommand(),
|
||||
freeCommand(),
|
||||
statusCommand(),
|
||||
checkCommand(),
|
||||
clearCommand(),
|
||||
subagentsCommand(),
|
||||
|
|
|
|||
144
pkg/commands/cmd_mode.go
Normal file
144
pkg/commands/cmd_mode.go
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
package commands
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func boostCommand() Definition {
|
||||
return Definition{
|
||||
Name: "boost",
|
||||
Description: "Use the paid model for your next message",
|
||||
Usage: "/boost",
|
||||
Handler: func(_ context.Context, req Request, rt *Runtime) error {
|
||||
paidModel, _ := sessionModelNames(rt)
|
||||
if paidModel == "" {
|
||||
return req.Reply("Boost unavailable: paid model is not configured.")
|
||||
}
|
||||
if rt == nil || rt.ArmNextModelMode == nil {
|
||||
return req.Reply(unavailableMsg)
|
||||
}
|
||||
if err := rt.ArmNextModelMode(paidModel); err != nil {
|
||||
return req.Reply(err.Error())
|
||||
}
|
||||
return req.Reply(fmt.Sprintf("Boost armed. Next message will use %s.", paidModel))
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func paidCommand() Definition {
|
||||
return Definition{
|
||||
Name: "paid",
|
||||
Description: "Use the paid model for this session",
|
||||
Usage: "/paid",
|
||||
Handler: func(_ context.Context, req Request, rt *Runtime) error {
|
||||
paidModel, _ := sessionModelNames(rt)
|
||||
if paidModel == "" {
|
||||
return req.Reply("Paid mode unavailable: primary model is not configured.")
|
||||
}
|
||||
if rt == nil || rt.SetSessionModelMode == nil {
|
||||
return req.Reply(unavailableMsg)
|
||||
}
|
||||
if err := rt.SetSessionModelMode(paidModel); err != nil {
|
||||
return req.Reply(err.Error())
|
||||
}
|
||||
return req.Reply(fmt.Sprintf("Session mode set to paid (%s).", paidModel))
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func freeCommand() Definition {
|
||||
return Definition{
|
||||
Name: "free",
|
||||
Description: "Use the free model for this session",
|
||||
Usage: "/free",
|
||||
Handler: func(_ context.Context, req Request, rt *Runtime) error {
|
||||
_, freeModel := sessionModelNames(rt)
|
||||
if freeModel == "" {
|
||||
return req.Reply("Free mode unavailable: light model is not configured.")
|
||||
}
|
||||
if rt == nil || rt.SetSessionModelMode == nil {
|
||||
return req.Reply(unavailableMsg)
|
||||
}
|
||||
if err := rt.SetSessionModelMode(freeModel); err != nil {
|
||||
return req.Reply(err.Error())
|
||||
}
|
||||
return req.Reply(fmt.Sprintf("Session mode set to free (%s).", freeModel))
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func statusCommand() Definition {
|
||||
return Definition{
|
||||
Name: "status",
|
||||
Description: "Show the current session model mode",
|
||||
Usage: "/status",
|
||||
Handler: func(_ context.Context, req Request, rt *Runtime) error {
|
||||
if rt == nil {
|
||||
return req.Reply(unavailableMsg)
|
||||
}
|
||||
|
||||
currentModel, provider := "", ""
|
||||
if rt.GetModelInfo != nil {
|
||||
currentModel, provider = rt.GetModelInfo()
|
||||
}
|
||||
|
||||
paidModel, freeModel := sessionModelNames(rt)
|
||||
persistent, pending := "", ""
|
||||
if rt.GetSessionModelMode != nil {
|
||||
persistent, pending = rt.GetSessionModelMode()
|
||||
}
|
||||
|
||||
lines := make([]string, 0, 5)
|
||||
if currentModel != "" {
|
||||
if provider != "" {
|
||||
lines = append(lines, fmt.Sprintf("Current Model: %s (Provider: %s)", currentModel, provider))
|
||||
} else {
|
||||
lines = append(lines, fmt.Sprintf("Current Model: %s", currentModel))
|
||||
}
|
||||
}
|
||||
lines = append(lines, fmt.Sprintf("Session Mode: %s", sessionModeDescription(persistent, pending, paidModel, freeModel)))
|
||||
if pending != "" {
|
||||
lines = append(lines, fmt.Sprintf("Pending Boost: %s", pending))
|
||||
} else {
|
||||
lines = append(lines, "Pending Boost: none")
|
||||
}
|
||||
if paidModel != "" {
|
||||
lines = append(lines, fmt.Sprintf("Paid Model: %s", paidModel))
|
||||
}
|
||||
if freeModel != "" {
|
||||
lines = append(lines, fmt.Sprintf("Free Model: %s", freeModel))
|
||||
}
|
||||
return req.Reply(strings.Join(lines, "\n"))
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func sessionModelNames(rt *Runtime) (paidModel, freeModel string) {
|
||||
if rt == nil || rt.Config == nil {
|
||||
return "", ""
|
||||
}
|
||||
|
||||
paidModel = strings.TrimSpace(rt.Config.Agents.Defaults.ModelName)
|
||||
if rt.Config.Agents.Defaults.Routing != nil {
|
||||
freeModel = strings.TrimSpace(rt.Config.Agents.Defaults.Routing.LightModel)
|
||||
}
|
||||
return paidModel, freeModel
|
||||
}
|
||||
|
||||
func sessionModeDescription(persistent, pending, paidModel, freeModel string) string {
|
||||
if pending != "" {
|
||||
return fmt.Sprintf("boost armed for next message (%s)", pending)
|
||||
}
|
||||
if persistent == "" {
|
||||
return "route (default)"
|
||||
}
|
||||
if paidModel != "" && strings.EqualFold(persistent, paidModel) {
|
||||
return fmt.Sprintf("paid (%s)", persistent)
|
||||
}
|
||||
if freeModel != "" && strings.EqualFold(persistent, freeModel) {
|
||||
return fmt.Sprintf("free (%s)", persistent)
|
||||
}
|
||||
return fmt.Sprintf("custom (%s)", persistent)
|
||||
}
|
||||
153
pkg/commands/cmd_mode_test.go
Normal file
153
pkg/commands/cmd_mode_test.go
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
package commands
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
)
|
||||
|
||||
func newModeTestRuntime() *Runtime {
|
||||
return &Runtime{
|
||||
Config: &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
ModelName: "gpt-5.4-mini",
|
||||
Routing: &config.RoutingConfig{
|
||||
LightModel: "openrouter-free",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestBoostCommand_ArmsNextModel(t *testing.T) {
|
||||
rt := newModeTestRuntime()
|
||||
var armed string
|
||||
rt.ArmNextModelMode = func(value string) error {
|
||||
armed = value
|
||||
return nil
|
||||
}
|
||||
|
||||
ex := NewExecutor(NewRegistry(BuiltinDefinitions()), rt)
|
||||
|
||||
var reply string
|
||||
res := ex.Execute(context.Background(), Request{
|
||||
Text: "/boost",
|
||||
Reply: func(text string) error {
|
||||
reply = text
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if res.Outcome != OutcomeHandled {
|
||||
t.Fatalf("outcome=%v, want=%v", res.Outcome, OutcomeHandled)
|
||||
}
|
||||
if armed != "gpt-5.4-mini" {
|
||||
t.Fatalf("armed=%q, want %q", armed, "gpt-5.4-mini")
|
||||
}
|
||||
if reply != "Boost armed. Next message will use gpt-5.4-mini." {
|
||||
t.Fatalf("reply=%q, want boost confirmation", reply)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPaidCommand_SetsPersistentModel(t *testing.T) {
|
||||
rt := newModeTestRuntime()
|
||||
var persistent string
|
||||
rt.SetSessionModelMode = func(value string) error {
|
||||
persistent = value
|
||||
return nil
|
||||
}
|
||||
|
||||
ex := NewExecutor(NewRegistry(BuiltinDefinitions()), rt)
|
||||
|
||||
var reply string
|
||||
res := ex.Execute(context.Background(), Request{
|
||||
Text: "/paid",
|
||||
Reply: func(text string) error {
|
||||
reply = text
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if res.Outcome != OutcomeHandled {
|
||||
t.Fatalf("outcome=%v, want=%v", res.Outcome, OutcomeHandled)
|
||||
}
|
||||
if persistent != "gpt-5.4-mini" {
|
||||
t.Fatalf("persistent=%q, want %q", persistent, "gpt-5.4-mini")
|
||||
}
|
||||
if reply != "Session mode set to paid (gpt-5.4-mini)." {
|
||||
t.Fatalf("reply=%q, want paid confirmation", reply)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFreeCommand_SetsPersistentModel(t *testing.T) {
|
||||
rt := newModeTestRuntime()
|
||||
var persistent string
|
||||
rt.SetSessionModelMode = func(value string) error {
|
||||
persistent = value
|
||||
return nil
|
||||
}
|
||||
|
||||
ex := NewExecutor(NewRegistry(BuiltinDefinitions()), rt)
|
||||
|
||||
var reply string
|
||||
res := ex.Execute(context.Background(), Request{
|
||||
Text: "/free",
|
||||
Reply: func(text string) error {
|
||||
reply = text
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if res.Outcome != OutcomeHandled {
|
||||
t.Fatalf("outcome=%v, want=%v", res.Outcome, OutcomeHandled)
|
||||
}
|
||||
if persistent != "openrouter-free" {
|
||||
t.Fatalf("persistent=%q, want %q", persistent, "openrouter-free")
|
||||
}
|
||||
if reply != "Session mode set to free (openrouter-free)." {
|
||||
t.Fatalf("reply=%q, want free confirmation", reply)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusCommand_ReportsPendingBoost(t *testing.T) {
|
||||
rt := newModeTestRuntime()
|
||||
rt.GetModelInfo = func() (string, string) {
|
||||
return "gpt-5.4-mini", "openai"
|
||||
}
|
||||
rt.GetSessionModelMode = func() (string, string) {
|
||||
return "openrouter-free", "gpt-5.4-mini"
|
||||
}
|
||||
|
||||
ex := NewExecutor(NewRegistry(BuiltinDefinitions()), rt)
|
||||
|
||||
var reply string
|
||||
res := ex.Execute(context.Background(), Request{
|
||||
Text: "/status",
|
||||
Reply: func(text string) error {
|
||||
reply = text
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if res.Outcome != OutcomeHandled {
|
||||
t.Fatalf("outcome=%v, want=%v", res.Outcome, OutcomeHandled)
|
||||
}
|
||||
if !containsAll(reply, []string{
|
||||
"Current Model: gpt-5.4-mini (Provider: openai)",
|
||||
"Session Mode: boost armed for next message (gpt-5.4-mini)",
|
||||
"Pending Boost: gpt-5.4-mini",
|
||||
"Paid Model: gpt-5.4-mini",
|
||||
"Free Model: openrouter-free",
|
||||
}) {
|
||||
t.Fatalf("reply=%q, missing expected status content", reply)
|
||||
}
|
||||
}
|
||||
|
||||
func containsAll(text string, want []string) bool {
|
||||
for _, s := range want {
|
||||
if !strings.Contains(text, s) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
|
@ -17,4 +17,8 @@ type Runtime struct {
|
|||
SwitchChannel func(value string) error
|
||||
ClearHistory func() error
|
||||
ReloadConfig func() error
|
||||
|
||||
GetSessionModelMode func() (persistent, pending string)
|
||||
SetSessionModelMode func(value string) error
|
||||
ArmNextModelMode func(value string) error
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue