opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix(agent): route image messages through image_model

Browse source
This commit is contained in:
Alix-007 2026-03-17 11:54:17 +08:00
parent 4a8a2e9c23
commit d8b4f088cb
5 changed files with 267 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
pkg/agent/instance.go
View file

@ -39,6 +39,8 @@ type AgentInstance struct {
Subagents *config.SubagentsConfig
SkillsFilter []string
Candidates []providers.FallbackCandidate
ImageModel string
ImageCandidates []providers.FallbackCandidate
// Router is non-nil when model routing is configured and the light model
// was successfully resolved. It scores each incoming message and decides
@ -195,6 +197,20 @@ func NewAgentInstance(
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
imageModel := strings.TrimSpace(defaults.ImageModel)
var imageCandidates []providers.FallbackCandidate
if imageModel != "" {
imageModelCfg := providers.ModelConfig{
Primary: imageModel,
Fallbacks: defaults.ImageModelFallbacks,
}
imageCandidates = providers.ResolveCandidatesWithLookup(
imageModelCfg,
defaults.Provider,
resolveFromModelList,
)
}
// Model routing setup: pre-resolve light model candidates at creation time
// to avoid repeated model_list lookups on every incoming message.
var router *routing.Router
@ -234,6 +250,8 @@ func NewAgentInstance(
Subagents: subagents,
SkillsFilter: skillsFilter,
Candidates: candidates,
ImageModel: imageModel,
ImageCandidates: imageCandidates,
Router: router,
LightCandidates: lightCandidates,
}

129
pkg/agent/instance_test.go
View file

@ -1,10 +1,14 @@
package agent
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/media"
)
func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) {
@ -160,3 +164,128 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) {
})
}
}
func TestNewAgentInstance_ResolveImageCandidatesFromModelListAlias(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "agent-instance-test-*")
if err != nil {
t.Fatalf("Failed to create temp dir: %v", err)
}
defer os.RemoveAll(tmpDir)
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: tmpDir,
Model: "text-main",
ImageModel: "vision-main",
ImageModelFallbacks: []string{"vision-backup"},
},
},
ModelList: []config.ModelConfig{
{
ModelName: "vision-main",
Model: "gemini/gemini-2.5-flash-lite",
},
{
ModelName: "vision-backup",
Model: "anthropic/claude-3-7-sonnet",
},
},
}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
if agent.ImageModel != "vision-main" {
t.Fatalf("ImageModel = %q, want %q", agent.ImageModel, "vision-main")
}
if len(agent.ImageCandidates) != 2 {
t.Fatalf("len(ImageCandidates) = %d, want 2", len(agent.ImageCandidates))
}
if agent.ImageCandidates[0].Provider != "gemini" || agent.ImageCandidates[0].Model != "gemini-2.5-flash-lite" {
t.Fatalf("first image candidate = %+v, want gemini/gemini-2.5-flash-lite", agent.ImageCandidates[0])
}
if agent.ImageCandidates[1].Provider != "anthropic" || agent.ImageCandidates[1].Model != "claude-3-7-sonnet" {
t.Fatalf("second image candidate = %+v, want anthropic/claude-3-7-sonnet", agent.ImageCandidates[1])
}
}
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
workspace := t.TempDir()
mediaDir := media.TempDir()
if err := os.MkdirAll(mediaDir, 0o700); err != nil {
t.Fatalf("MkdirAll(mediaDir) error = %v", err)
}
mediaFile, err := os.CreateTemp(mediaDir, "instance-tool-*.txt")
if err != nil {
t.Fatalf("CreateTemp(mediaDir) error = %v", err)
}
mediaPath := mediaFile.Name()
if _, err := mediaFile.WriteString("attachment content"); err != nil {
mediaFile.Close()
t.Fatalf("WriteString(mediaFile) error = %v", err)
}
if err := mediaFile.Close(); err != nil {
t.Fatalf("Close(mediaFile) error = %v", err)
}
t.Cleanup(func() { _ = os.Remove(mediaPath) })
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: workspace,
Model: "test-model",
RestrictToWorkspace: true,
},
},
Tools: config.ToolsConfig{
ReadFile: config.ReadFileToolConfig{Enabled: true},
ListDir: config.ToolConfig{Enabled: true},
Exec: config.ExecConfig{
ToolConfig: config.ToolConfig{Enabled: true},
EnableDenyPatterns: true,
AllowRemote: true,
},
},
}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
readTool, ok := agent.Tools.Get("read_file")
if !ok {
t.Fatal("read_file tool not registered")
}
readResult := readTool.Execute(context.Background(), map[string]any{"path": mediaPath})
if readResult.IsError {
t.Fatalf("read_file should allow media temp dir, got: %s", readResult.ForLLM)
}
if !strings.Contains(readResult.ForLLM, "attachment content") {
t.Fatalf("read_file output missing media content: %s", readResult.ForLLM)
}
listTool, ok := agent.Tools.Get("list_dir")
if !ok {
t.Fatal("list_dir tool not registered")
}
listResult := listTool.Execute(context.Background(), map[string]any{"path": mediaDir})
if listResult.IsError {
t.Fatalf("list_dir should allow media temp dir, got: %s", listResult.ForLLM)
}
if !strings.Contains(listResult.ForLLM, filepath.Base(mediaPath)) {
t.Fatalf("list_dir output missing media file: %s", listResult.ForLLM)
}
execTool, ok := agent.Tools.Get("exec")
if !ok {
t.Fatal("exec tool not registered")
}
execResult := execTool.Execute(context.Background(), map[string]any{
"command": "cat " + filepath.Base(mediaPath),
"working_dir": mediaDir,
})
if execResult.IsError {
t.Fatalf("exec should allow media temp dir, got: %s", execResult.ForLLM)
}
if !strings.Contains(execResult.ForLLM, "attachment content") {
t.Fatalf("exec output missing media content: %s", execResult.ForLLM)
}
}

48
pkg/agent/loop.go
View file

@ -982,7 +982,7 @@ func (al *AgentLoop) runLLMIteration(
// selectCandidates evaluates routing once and the decision is sticky for
// all tool-follow-up iterations within the same turn so that a multi-step
// tool chain doesn't switch models mid-way through.
activeCandidates, activeModel := al.selectCandidates(agent, opts.UserMessage, messages)
activeCandidates, activeModel, useImageFallback := al.selectCandidates(agent, opts.UserMessage, messages)
for iteration < agent.MaxIterations {
iteration++
@ -1040,13 +1040,19 @@ func (al *AgentLoop) runLLMIteration(
callLLM := func() (*providers.LLMResponse, error) {
if len(activeCandidates) > 1 && al.fallback != nil {
fbResult, fbErr := al.fallback.Execute(
ctx,
activeCandidates,
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
},
runCandidate := func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
}
var (
fbResult *providers.FallbackResult
fbErr error
)
if useImageFallback {
fbResult, fbErr = al.fallback.ExecuteImage(ctx, activeCandidates, runCandidate)
} else {
fbResult, fbErr = al.fallback.Execute(ctx, activeCandidates, runCandidate)
}
if fbErr != nil {
return nil, fbErr
}
@ -1387,9 +1393,18 @@ func (al *AgentLoop) selectCandidates(
agent *AgentInstance,
userMsg string,
history []providers.Message,
) (candidates []providers.FallbackCandidate, model string) {
) (candidates []providers.FallbackCandidate, model string, useImageFallback bool) {
if hasImageMedia(history) && len(agent.ImageCandidates) > 0 {
logger.InfoCF("agent", "Image model selected",
map[string]any{
"agent_id": agent.ID,
"image_model": agent.ImageModel,
})
return agent.ImageCandidates, agent.ImageModel, true
}
if agent.Router == nil || len(agent.LightCandidates) == 0 {
return agent.Candidates, agent.Model
return agent.Candidates, agent.Model, false
}
_, usedLight, score := agent.Router.SelectModel(userMsg, history, agent.Model)
@ -1400,7 +1415,7 @@ func (al *AgentLoop) selectCandidates(
"score": score,
"threshold": agent.Router.Threshold(),
})
return agent.Candidates, agent.Model
return agent.Candidates, agent.Model, false
}
logger.InfoCF("agent", "Model routing: light model selected",
@ -1410,7 +1425,18 @@ func (al *AgentLoop) selectCandidates(
"score": score,
"threshold": agent.Router.Threshold(),
})
return agent.LightCandidates, agent.Router.LightModel()
return agent.LightCandidates, agent.Router.LightModel(), false
}
func hasImageMedia(messages []providers.Message) bool {
for _, msg := range messages {
for _, ref := range msg.Media {
if strings.HasPrefix(strings.ToLower(ref), "data:image/") {
return true
}
}
}
return false
}
// maybeSummarize triggers summarization if the session history exceeds thresholds.

79
pkg/agent/loop_test.go
View file

@ -88,6 +88,85 @@ func TestRecordLastChatID(t *testing.T) {
}
}
func TestSelectCandidates_UsesImageModelWhenImageMediaPresent(t *testing.T) {
tmpDir := t.TempDir()
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: tmpDir,
ModelName: "text-main",
ImageModel: "vision-main",
},
},
ModelList: []config.ModelConfig{
{ModelName: "text-main", Model: "openai/gpt-5.4"},
{ModelName: "vision-main", Model: "gemini/gemini-2.5-flash-lite"},
},
}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
candidates, model, useImageFallback := (&AgentLoop{}).selectCandidates(
agent,
"describe this image",
[]providers.Message{
{Role: "user", Content: "describe this image", Media: []string{"data:image/png;base64,AAAA"}},
},
)
if !useImageFallback {
t.Fatal("expected image fallback to be selected")
}
if model != "vision-main" {
t.Fatalf("model = %q, want %q", model, "vision-main")
}
if len(candidates) != 1 {
t.Fatalf("len(candidates) = %d, want 1", len(candidates))
}
if candidates[0].Provider != "gemini" || candidates[0].Model != "gemini-2.5-flash-lite" {
t.Fatalf("candidate = %+v, want gemini/gemini-2.5-flash-lite", candidates[0])
}
}
func TestRunAgentLoop_UsesImageModelForImageMessages(t *testing.T) {
tmpDir := t.TempDir()
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: tmpDir,
ModelName: "text-main",
ImageModel: "vision-main",
MaxTokens: 4096,
MaxToolIterations: 5,
},
},
ModelList: []config.ModelConfig{
{ModelName: "text-main", Model: "openai/gpt-5.4"},
{ModelName: "vision-main", Model: "gemini/gemini-2.5-flash-lite"},
},
}
msgBus := bus.NewMessageBus()
provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider)
agent := al.registry.GetDefaultAgent()
_, err := al.runAgentLoop(context.Background(), agent, processOptions{
SessionKey: "image-session",
Channel: "telegram",
ChatID: "chat-1",
UserMessage: "describe this image",
Media: []string{"data:image/png;base64,AAAA"},
DefaultResponse: "fallback",
SendResponse: false,
})
if err != nil {
t.Fatalf("runAgentLoop returned error: %v", err)
}
if provider.lastModel != "vision-main" {
t.Fatalf("provider lastModel = %q, want %q", provider.lastModel, "vision-main")
}
}
func TestNewAgentLoop_StateInitialized(t *testing.T) {
// Create temp workspace
tmpDir, err := os.MkdirTemp("", "agent-test-*")

5
pkg/agent/mock_provider_test.go
View file

@ -6,7 +6,9 @@ import (
"github.com/sipeed/picoclaw/pkg/providers"
)
type mockProvider struct{}
type mockProvider struct {
lastModel string
}
func (m *mockProvider) Chat(
ctx context.Context,
@ -15,6 +17,7 @@ func (m *mockProvider) Chat(
model string,
opts map[string]any,
) (*providers.LLMResponse, error) {
m.lastModel = model
return &providers.LLMResponse{
Content: "Mock response",
ToolCalls: []providers.ToolCall{},