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
|
|
@ -47,18 +47,20 @@ type AgentLoop struct {
|
||||||
hooks *HookManager
|
hooks *HookManager
|
||||||
|
|
||||||
// Runtime state
|
// Runtime state
|
||||||
running atomic.Bool
|
running atomic.Bool
|
||||||
contextManager ContextManager
|
contextManager ContextManager
|
||||||
fallback *providers.FallbackChain
|
fallback *providers.FallbackChain
|
||||||
channelManager *channels.Manager
|
channelManager *channels.Manager
|
||||||
mediaStore media.MediaStore
|
mediaStore media.MediaStore
|
||||||
transcriber asr.Transcriber
|
transcriber asr.Transcriber
|
||||||
cmdRegistry *commands.Registry
|
cmdRegistry *commands.Registry
|
||||||
mcp mcpRuntime
|
mcp mcpRuntime
|
||||||
hookRuntime hookRuntime
|
hookRuntime hookRuntime
|
||||||
steering *steeringQueue
|
steering *steeringQueue
|
||||||
pendingSkills sync.Map
|
pendingSkills sync.Map
|
||||||
mu sync.RWMutex
|
pendingModelOverrides sync.Map
|
||||||
|
sessionModelOverrides sync.Map
|
||||||
|
mu sync.RWMutex
|
||||||
|
|
||||||
// Concurrent turn management (from HEAD)
|
// Concurrent turn management (from HEAD)
|
||||||
activeTurnStates sync.Map // key: sessionKey (string), value: *turnState
|
activeTurnStates sync.Map // key: sessionKey (string), value: *turnState
|
||||||
|
|
@ -82,6 +84,7 @@ type processOptions struct {
|
||||||
SenderDisplayName string // Current sender display name for dynamic context
|
SenderDisplayName string // Current sender display name for dynamic context
|
||||||
UserMessage string // User message content (may include prefix)
|
UserMessage string // User message content (may include prefix)
|
||||||
ForcedSkills []string // Skills explicitly requested for this message
|
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)
|
SystemPromptOverride string // Override the default system prompt (Used by SubTurns)
|
||||||
Media []string // media:// refs from inbound message
|
Media []string // media:// refs from inbound message
|
||||||
InitialSteeringMessages []providers.Message // Steering messages from refactor/agent
|
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
|
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 {
|
if pending := al.takePendingSkills(opts.SessionKey); len(pending) > 0 {
|
||||||
opts.ForcedSkills = append(opts.ForcedSkills, pending...)
|
opts.ForcedSkills = append(opts.ForcedSkills, pending...)
|
||||||
logger.InfoCF("agent", "Applying pending skill override",
|
logger.InfoCF("agent", "Applying pending skill override",
|
||||||
|
|
@ -1784,10 +1803,34 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
|
||||||
ts.ingestMessage(turnCtx, al, rootMsg)
|
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
|
activeProvider := ts.agent.Provider
|
||||||
if usedLight && ts.agent.LightProvider != nil {
|
activeThinkingLevel := ts.agent.ThinkingLevel
|
||||||
activeProvider = ts.agent.LightProvider
|
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...)
|
pendingMessages := append([]providers.Message(nil), ts.opts.InitialSteeringMessages...)
|
||||||
var finalContent string
|
var finalContent string
|
||||||
|
|
@ -1893,7 +1936,7 @@ turnLoop:
|
||||||
hasWebSearch &&
|
hasWebSearch &&
|
||||||
func() bool {
|
func() bool {
|
||||||
// Check if provider supports native search
|
// 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 ns.SupportsNativeSearch()
|
||||||
}
|
}
|
||||||
return false
|
return false
|
||||||
|
|
@ -1933,12 +1976,12 @@ turnLoop:
|
||||||
if useNativeSearch {
|
if useNativeSearch {
|
||||||
llmOpts["native_search"] = true
|
llmOpts["native_search"] = true
|
||||||
}
|
}
|
||||||
if ts.agent.ThinkingLevel != ThinkingOff {
|
if activeThinkingLevel != ThinkingOff {
|
||||||
if tc, ok := ts.agent.Provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() {
|
if tc, ok := activeProvider.(providers.ThinkingCapable); ok && tc.SupportsThinking() {
|
||||||
llmOpts["thinking_level"] = string(ts.agent.ThinkingLevel)
|
llmOpts["thinking_level"] = string(activeThinkingLevel)
|
||||||
} else {
|
} else {
|
||||||
logger.WarnCF("agent", "thinking_level is set but current provider does not support it, ignoring",
|
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, ""
|
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 {
|
func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOptions) *commands.Runtime {
|
||||||
registry := al.GetRegistry()
|
registry := al.GetRegistry()
|
||||||
cfg := al.GetConfig()
|
cfg := al.GetConfig()
|
||||||
|
|
@ -3265,6 +3337,42 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
||||||
return oldModel, nil
|
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 {
|
rt.ClearHistory = func() error {
|
||||||
if opts == nil {
|
if opts == nil {
|
||||||
return fmt.Errorf("process options not available")
|
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) {
|
func (al *AgentLoop) setPendingSkills(sessionKey string, skillNames []string) {
|
||||||
sessionKey = strings.TrimSpace(sessionKey)
|
sessionKey = strings.TrimSpace(sessionKey)
|
||||||
if sessionKey == "" || len(skillNames) == 0 {
|
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) {
|
func TestProcessMessage_ModelRoutingUsesLightProvider(t *testing.T) {
|
||||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|
|
||||||
|
|
@ -12,6 +12,10 @@ func BuiltinDefinitions() []Definition {
|
||||||
listCommand(),
|
listCommand(),
|
||||||
useCommand(),
|
useCommand(),
|
||||||
switchCommand(),
|
switchCommand(),
|
||||||
|
boostCommand(),
|
||||||
|
paidCommand(),
|
||||||
|
freeCommand(),
|
||||||
|
statusCommand(),
|
||||||
checkCommand(),
|
checkCommand(),
|
||||||
clearCommand(),
|
clearCommand(),
|
||||||
subagentsCommand(),
|
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
|
SwitchChannel func(value string) error
|
||||||
ClearHistory func() error
|
ClearHistory func() error
|
||||||
ReloadConfig 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