This commit is contained in:
Anton Bogdanovich 2026-05-15 11:29:28 +03:00 committed by GitHub
commit 4347a4f1cf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 893 additions and 60 deletions

View file

@ -442,6 +442,10 @@
"i2c": {
"enabled": false
},
"image_generate": {
"enabled": false,
"model": "openai-codex/gpt-image-2"
},
"install_skill": {
"enabled": true
},

View file

@ -182,6 +182,30 @@ If `range` is omitted, PicoClaw performs an unrestricted search.
}
```
## Image Generation Tool
The `image_generate` tool creates image files through a provider that supports
image generation.
| Config | Type | Default | Description |
|--------|------|---------|-------------|
| `enabled` | bool | false | Enable the image generation tool |
| `model` | string | `gpt-image-2` | Image generation model. Values may include a provider prefix, for example `openai-codex/gpt-image-2` |
If `tools.image_generate.model` is not set, PicoClaw falls back to the legacy
`agents.defaults.image_model` setting, then to `gpt-image-2`.
```json
{
"tools": {
"image_generate": {
"enabled": true,
"model": "openai-codex/gpt-image-2"
}
}
}
```
## Exec Tool
The exec tool is used to execute shell commands.

View file

@ -231,6 +231,11 @@ func registerSharedTools(
agent.Tools.Register(loadImageTool)
}
if cfg.Tools.IsToolEnabled("image_generate") {
imageModel := cfg.Tools.ImageGenerate.EffectiveModel(cfg.Agents.Defaults)
agent.Tools.Register(tools.NewImageGenerateTool(agent.Workspace, imageModel, nil))
}
// Skill discovery and installation tools
skills_enabled := cfg.Tools.IsToolEnabled("skills")
find_skills_enable := cfg.Tools.IsToolEnabled("find_skills")

View file

@ -786,6 +786,21 @@ type ToolConfig struct {
Enabled bool `json:"enabled" yaml:"-" env:"ENABLED"`
}
type ImageGenerateToolsConfig struct {
ToolConfig `yaml:"-" envPrefix:"PICOCLAW_TOOLS_IMAGE_GENERATE_"`
Model string `json:"model,omitempty" yaml:"-" env:"PICOCLAW_TOOLS_IMAGE_GENERATE_MODEL"`
}
func (c ImageGenerateToolsConfig) EffectiveModel(defaults AgentDefaults) string {
if model := strings.TrimSpace(c.Model); model != "" {
return model
}
if model := strings.TrimSpace(defaults.ImageModel); model != "" {
return model
}
return "gpt-image-2"
}
type BraveConfig struct {
Enabled bool `json:"enabled" yaml:"-" env:"PICOCLAW_TOOLS_WEB_BRAVE_ENABLED"`
APIKeys SecureStrings `json:"api_keys,omitzero" yaml:"api_keys,omitempty" env:"PICOCLAW_TOOLS_WEB_BRAVE_API_KEYS"`
@ -984,30 +999,31 @@ type ToolsConfig struct {
// FilterMinLength is the minimum content length required for filtering.
// Content shorter than this will be returned unchanged for performance.
// Default: 8
FilterMinLength int `json:"filter_min_length" yaml:"-" env:"PICOCLAW_TOOLS_FILTER_MIN_LENGTH"`
Web WebToolsConfig `json:"web" yaml:"web,omitempty"`
Cron CronToolsConfig `json:"cron" yaml:"-"`
Exec ExecConfig `json:"exec" yaml:"-"`
Skills SkillsToolsConfig `json:"skills" yaml:"skills,omitempty"`
MediaCleanup MediaCleanupConfig `json:"media_cleanup" yaml:"-"`
MCP MCPConfig `json:"mcp" yaml:"-"`
AppendFile ToolConfig `json:"append_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_APPEND_FILE_"`
EditFile ToolConfig `json:"edit_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_EDIT_FILE_"`
FindSkills ToolConfig `json:"find_skills" yaml:"-" envPrefix:"PICOCLAW_TOOLS_FIND_SKILLS_"`
I2C ToolConfig `json:"i2c" yaml:"-" envPrefix:"PICOCLAW_TOOLS_I2C_"`
InstallSkill ToolConfig `json:"install_skill" yaml:"-" envPrefix:"PICOCLAW_TOOLS_INSTALL_SKILL_"`
ListDir ToolConfig `json:"list_dir" yaml:"-" envPrefix:"PICOCLAW_TOOLS_LIST_DIR_"`
Message ToolConfig `json:"message" yaml:"-" envPrefix:"PICOCLAW_TOOLS_MESSAGE_"`
ReadFile ReadFileToolConfig `json:"read_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_READ_FILE_"`
Serial ToolConfig `json:"serial" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SERIAL_"`
SendFile ToolConfig `json:"send_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SEND_FILE_"`
SendTTS ToolConfig `json:"send_tts" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SEND_TTS_"`
Spawn ToolConfig `json:"spawn" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SPAWN_"`
SpawnStatus ToolConfig `json:"spawn_status" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SPAWN_STATUS_"`
SPI ToolConfig `json:"spi" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SPI_"`
Subagent ToolConfig `json:"subagent" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SUBAGENT_"`
WebFetch ToolConfig `json:"web_fetch" yaml:"-" envPrefix:"PICOCLAW_TOOLS_WEB_FETCH_"`
WriteFile ToolConfig `json:"write_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_WRITE_FILE_"`
FilterMinLength int `json:"filter_min_length" yaml:"-" env:"PICOCLAW_TOOLS_FILTER_MIN_LENGTH"`
Web WebToolsConfig `json:"web" yaml:"web,omitempty"`
Cron CronToolsConfig `json:"cron" yaml:"-"`
Exec ExecConfig `json:"exec" yaml:"-"`
Skills SkillsToolsConfig `json:"skills" yaml:"skills,omitempty"`
MediaCleanup MediaCleanupConfig `json:"media_cleanup" yaml:"-"`
MCP MCPConfig `json:"mcp" yaml:"-"`
AppendFile ToolConfig `json:"append_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_APPEND_FILE_"`
EditFile ToolConfig `json:"edit_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_EDIT_FILE_"`
FindSkills ToolConfig `json:"find_skills" yaml:"-" envPrefix:"PICOCLAW_TOOLS_FIND_SKILLS_"`
I2C ToolConfig `json:"i2c" yaml:"-" envPrefix:"PICOCLAW_TOOLS_I2C_"`
ImageGenerate ImageGenerateToolsConfig `json:"image_generate" yaml:"-"`
InstallSkill ToolConfig `json:"install_skill" yaml:"-" envPrefix:"PICOCLAW_TOOLS_INSTALL_SKILL_"`
ListDir ToolConfig `json:"list_dir" yaml:"-" envPrefix:"PICOCLAW_TOOLS_LIST_DIR_"`
Message ToolConfig `json:"message" yaml:"-" envPrefix:"PICOCLAW_TOOLS_MESSAGE_"`
ReadFile ReadFileToolConfig `json:"read_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_READ_FILE_"`
Serial ToolConfig `json:"serial" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SERIAL_"`
SendFile ToolConfig `json:"send_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SEND_FILE_"`
SendTTS ToolConfig `json:"send_tts" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SEND_TTS_"`
Spawn ToolConfig `json:"spawn" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SPAWN_"`
SpawnStatus ToolConfig `json:"spawn_status" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SPAWN_STATUS_"`
SPI ToolConfig `json:"spi" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SPI_"`
Subagent ToolConfig `json:"subagent" yaml:"-" envPrefix:"PICOCLAW_TOOLS_SUBAGENT_"`
WebFetch ToolConfig `json:"web_fetch" yaml:"-" envPrefix:"PICOCLAW_TOOLS_WEB_FETCH_"`
WriteFile ToolConfig `json:"write_file" yaml:"-" envPrefix:"PICOCLAW_TOOLS_WRITE_FILE_"`
}
// IsFilterSensitiveDataEnabled returns true if sensitive data filtering is enabled
@ -1719,6 +1735,8 @@ func (t *ToolsConfig) IsToolEnabled(name string) bool {
return t.FindSkills.Enabled
case "i2c":
return t.I2C.Enabled
case "image_generate":
return t.ImageGenerate.Enabled
case "install_skill":
return t.InstallSkill.Enabled
case "list_dir":

View file

@ -506,6 +506,50 @@ func TestLoadConfig_MCPMaxInlineTextChars(t *testing.T) {
}
}
func TestImageGenerateToolsConfig_EffectiveModel(t *testing.T) {
defaults := AgentDefaults{ImageModel: "legacy-image-model"}
if got := (ImageGenerateToolsConfig{}).EffectiveModel(defaults); got != "legacy-image-model" {
t.Fatalf("legacy fallback model = %q, want legacy-image-model", got)
}
cfg := ImageGenerateToolsConfig{Model: "openai-codex/gpt-image-2"}
if got := cfg.EffectiveModel(defaults); got != "openai-codex/gpt-image-2" {
t.Fatalf("tool model = %q, want openai-codex/gpt-image-2", got)
}
if got := (ImageGenerateToolsConfig{}).EffectiveModel(AgentDefaults{}); got != "gpt-image-2" {
t.Fatalf("default model = %q, want gpt-image-2", got)
}
}
func TestLoadConfig_ImageGenerateModel(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "config.json")
raw := `{
"tools": {
"image_generate": {
"enabled": true,
"model": "openai-codex/gpt-image-2"
}
}
}`
if err := os.WriteFile(configPath, []byte(raw), 0o644); err != nil {
t.Fatalf("WriteFile(configPath): %v", err)
}
cfg, err := LoadConfig(configPath)
if err != nil {
t.Fatalf("LoadConfig() error: %v", err)
}
if !cfg.Tools.ImageGenerate.Enabled {
t.Fatal("cfg.Tools.ImageGenerate.Enabled should be true")
}
if got := cfg.Tools.ImageGenerate.Model; got != "openai-codex/gpt-image-2" {
t.Fatalf("cfg.Tools.ImageGenerate.Model = %q, want openai-codex/gpt-image-2", got)
}
}
func TestConfig_BackwardCompat_NoAgentsList(t *testing.T) {
jsonData := `{
"agents": {

View file

@ -409,6 +409,11 @@ func DefaultConfig() *Config {
SendFile: ToolConfig{
Enabled: true,
},
ImageGenerate: ImageGenerateToolsConfig{
ToolConfig: ToolConfig{
Enabled: false,
},
},
SendTTS: ToolConfig{
Enabled: false,
},

View file

@ -407,6 +407,30 @@ func CreateProviderFromConfig(cfg *config.ModelConfig) (LLMProvider, string, err
}
}
// CreateImageGenerationProviderFromModel creates a provider for image generation
// from a provider-prefixed model string. It returns the provider plus the model
// identifier stripped of the provider prefix.
func CreateImageGenerationProviderFromModel(model string) (ImageGenerationCapable, string, error) {
providerName, modelID := ExtractProtocol(&config.ModelConfig{Model: model})
if modelID == "" {
modelID = "gpt-image-2"
}
switch providerName {
case "", "openai", "openai-codex":
provider, err := createCodexAuthProvider()
if err != nil {
return nil, "", err
}
imageProvider, ok := provider.(ImageGenerationCapable)
if !ok || !imageProvider.SupportsImageGeneration() {
return nil, "", fmt.Errorf("provider %q does not support image generation", providerName)
}
return imageProvider, modelID, nil
default:
return nil, "", fmt.Errorf("provider %q does not support image generation", providerName)
}
}
func finalizeProviderFromConfig(
provider LLMProvider,
modelID string,

View file

@ -104,6 +104,32 @@ func TestCreateProviderReturnsClaudeProviderForAnthropicOAuth(t *testing.T) {
// TODO: Test custom APIBase when createClaudeAuthProvider supports it
}
func TestCreateImageGenerationProviderFromModelUsesCodexOAuth(t *testing.T) {
originalGetCredential := getCredential
t.Cleanup(func() { getCredential = originalGetCredential })
getCredential = func(provider string) (*auth.AuthCredential, error) {
if provider != "openai" {
t.Fatalf("provider = %q, want openai", provider)
}
return &auth.AuthCredential{
AccessToken: "openai-token",
AccountID: "acct-123",
}, nil
}
provider, model, err := CreateImageGenerationProviderFromModel("openai/gpt-image-2")
if err != nil {
t.Fatalf("CreateImageGenerationProviderFromModel() error = %v", err)
}
if model != "gpt-image-2" {
t.Fatalf("model = %q, want gpt-image-2", model)
}
if provider.ImageGenerationProviderID() != "openai-codex" {
t.Fatalf("provider id = %q, want openai-codex", provider.ImageGenerationProviderID())
}
}
func TestCreateProviderReturnsCodexProviderForOpenAIOAuth(t *testing.T) {
// TODO: This test requires openai protocol to support auth_method: "oauth"
// which is not yet implemented in the new factory_provider.go

View file

@ -0,0 +1,242 @@
package oauthprovider
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"strings"
"github.com/openai/openai-go/v3"
"github.com/openai/openai-go/v3/option"
"github.com/openai/openai-go/v3/responses"
)
const (
codexDefaultImageGenerationModel = "gpt-image-2"
codexDefaultImageGenerationSize = "1024x1024"
maxImageGenerationResults = 4
maxImageGenerationSSEBytes = 64 * 1024 * 1024
maxImageGenerationEvents = 512
)
func (p *CodexProvider) SupportsImageGeneration() bool {
return true
}
func (p *CodexProvider) ImageGenerationProviderID() string {
return "openai-codex"
}
func (p *CodexProvider) DefaultImageGenerationModel() string {
return codexDefaultImageGenerationModel
}
func (p *CodexProvider) GenerateImage(
ctx context.Context,
req ImageGenerationRequest,
) (*ImageGenerationResponse, error) {
opts, accountID, err := p.requestOptions()
if err != nil {
return nil, err
}
if accountID == "" {
return nil, fmt.Errorf("no account id found for Codex image generation")
}
if strings.TrimSpace(req.Model) == "" {
req.Model = p.DefaultImageGenerationModel()
}
if req.Count < 1 {
req.Count = 1
}
if req.Count > maxImageGenerationResults {
req.Count = maxImageGenerationResults
}
images := make([]GeneratedImage, 0, req.Count)
for i := 0; i < req.Count; i++ {
params := buildCodexImageParams(req)
stream := p.client.Responses.NewStreaming(ctx, params, opts...)
eventImages, readErr := parseCodexImageSSE(stream, req.OutputFormat)
closeErr := stream.Close()
if readErr != nil {
return nil, readErr
}
if closeErr != nil {
return nil, closeErr
}
images = append(images, eventImages...)
}
if len(images) > maxImageGenerationResults {
images = images[:maxImageGenerationResults]
}
return &ImageGenerationResponse{Images: images}, nil
}
func (p *CodexProvider) requestOptions() ([]option.RequestOption, string, error) {
var opts []option.RequestOption
accountID := p.accountID
if p.tokenSource != nil {
tok, accID, err := p.tokenSource()
if err != nil {
return nil, "", fmt.Errorf("refreshing token: %w", err)
}
opts = append(opts, option.WithAPIKey(tok))
if accID != "" {
accountID = accID
}
}
if accountID != "" {
opts = append(opts, option.WithHeader("Chatgpt-Account-Id", accountID))
}
return opts, accountID, nil
}
func buildCodexImageParams(req ImageGenerationRequest) responses.ResponseNewParams {
size := strings.TrimSpace(req.Size)
if size == "" {
size = codexDefaultImageGenerationSize
}
tool := responses.ToolUnionParam{OfImageGeneration: &responses.ToolImageGenerationParam{
Model: req.Model,
Size: size,
}}
if req.Quality != "" {
tool.OfImageGeneration.Quality = req.Quality
}
if req.OutputFormat != "" {
tool.OfImageGeneration.OutputFormat = req.OutputFormat
}
content := responses.ResponseInputMessageContentListParam{
responses.ResponseInputContentParamOfInputText(req.Prompt),
}
input := responses.ResponseInputParam{
responses.ResponseInputItemParamOfMessage(content, responses.EasyInputMessageRoleUser),
}
return responses.ResponseNewParams{
Model: "gpt-5.4",
Input: responses.ResponseNewParamsInputUnion{
OfInputItemList: input,
},
Instructions: openai.Opt("You are an image generation assistant."),
Tools: []responses.ToolUnionParam{tool},
ToolChoice: responses.ResponseNewParamsToolChoiceUnion{
OfHostedTool: &responses.ToolChoiceTypesParam{
Type: responses.ToolChoiceTypesTypeImageGeneration,
},
},
Store: openai.Opt(false),
}
}
type codexImageStream interface {
Next() bool
Current() responses.ResponseStreamEventUnion
Err() error
}
func parseCodexImageSSE(stream codexImageStream, outputFormat string) ([]GeneratedImage, error) {
var totalBytes int
var events int
var images []GeneratedImage
var completedImages []GeneratedImage
for stream.Next() {
evt := stream.Current()
events++
if events > maxImageGenerationEvents {
return nil, fmt.Errorf("codex image response exceeded event limit")
}
data, err := json.Marshal(evt)
if err == nil {
totalBytes += len(data)
if totalBytes > maxImageGenerationSSEBytes {
return nil, fmt.Errorf("codex image response exceeded size limit")
}
}
eventImages, eventCompletedImages, parseErr := parseCodexImageEventUnion(evt, outputFormat)
if parseErr != nil {
return nil, parseErr
}
images = append(images, eventImages...)
completedImages = append(completedImages, eventCompletedImages...)
}
if err := stream.Err(); err != nil {
return nil, err
}
if len(images) > 0 {
return images, nil
}
return completedImages, nil
}
func parseCodexImageEventUnion(
evt responses.ResponseStreamEventUnion,
outputFormat string,
) ([]GeneratedImage, []GeneratedImage, error) {
switch evt.Type {
case "response.output_item.done":
if image, ok, err := imageFromCodexItemUnion(evt.Item, outputFormat); err != nil {
return nil, nil, err
} else if ok {
return []GeneratedImage{image}, nil, nil
}
case "response.completed":
images := make([]GeneratedImage, 0)
for _, item := range evt.Response.Output {
if image, ok, err := imageFromCodexResponseItem(item, outputFormat); err != nil {
return nil, nil, err
} else if ok {
images = append(images, image)
}
}
return nil, images, nil
case "response.failed", "error":
return nil, nil, fmt.Errorf("codex image generation failed")
}
return nil, nil, nil
}
func imageFromCodexItemUnion(
item responses.ResponseOutputItemUnion,
outputFormat string,
) (GeneratedImage, bool, error) {
if item.Type != "image_generation_call" {
return GeneratedImage{}, false, nil
}
return imageFromCodexPayload(item.Result, outputFormat)
}
func imageFromCodexResponseItem(
item responses.ResponseOutputItemUnion,
outputFormat string,
) (GeneratedImage, bool, error) {
return imageFromCodexItemUnion(item, outputFormat)
}
func imageFromCodexPayload(payload string, outputFormat string) (GeneratedImage, bool, error) {
if payload == "" {
return GeneratedImage{}, false, nil
}
data, err := base64.StdEncoding.DecodeString(payload)
if err != nil {
return GeneratedImage{}, false, err
}
mime, ext := imageMimeAndExtension(outputFormat)
return GeneratedImage{Data: data, MimeType: mime, Ext: ext}, true, nil
}
func imageMimeAndExtension(outputFormat string) (string, string) {
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
case "jpeg", "jpg":
return "image/jpeg", "jpg"
case "webp":
return "image/webp", "webp"
default:
return "image/png", "png"
}
}

View file

@ -0,0 +1,111 @@
package oauthprovider
import (
"encoding/base64"
"encoding/json"
"strings"
"testing"
"github.com/openai/openai-go/v3/responses"
)
type mockCodexImageStream struct {
events []responses.ResponseStreamEventUnion
index int
err error
}
func (s *mockCodexImageStream) Next() bool {
if s.index >= len(s.events) {
return false
}
s.index++
return true
}
func (s *mockCodexImageStream) Current() responses.ResponseStreamEventUnion {
return s.events[s.index-1]
}
func (s *mockCodexImageStream) Err() error { return s.err }
func TestCodexProviderSupportsImageGeneration(t *testing.T) {
provider := NewCodexProvider("test-token", "acct-123")
if !provider.SupportsImageGeneration() {
t.Fatal("SupportsImageGeneration = false, want true")
}
if provider.ImageGenerationProviderID() != "openai-codex" {
t.Fatalf("provider id = %q, want openai-codex", provider.ImageGenerationProviderID())
}
if provider.DefaultImageGenerationModel() != "gpt-image-2" {
t.Fatalf("default image model = %q, want gpt-image-2", provider.DefaultImageGenerationModel())
}
}
func TestBuildCodexImageParams(t *testing.T) {
params := buildCodexImageParams(ImageGenerationRequest{
Prompt: "make a tiny icon",
Model: "gpt-image-2",
Size: "1536x1024",
Quality: "medium",
OutputFormat: "png",
})
if params.Model != "gpt-5.4" {
t.Fatalf("request model = %q, want gpt-5.4", params.Model)
}
if params.Input.OfString.Valid() {
t.Fatalf("input uses string form, want structured message input")
}
if len(params.Input.OfInputItemList) != 1 {
t.Fatalf("input item count = %d, want 1", len(params.Input.OfInputItemList))
}
data, err := json.Marshal(params)
if err != nil {
t.Fatalf("marshal params: %v", err)
}
payload := string(data)
for _, want := range []string{`"input":[`, `"role":"user"`, `"type":"input_text"`, `"text":"make a tiny icon"`} {
if !strings.Contains(payload, want) {
t.Fatalf("payload missing %s: %s", want, payload)
}
}
if len(params.Tools) != 1 || params.Tools[0].OfImageGeneration == nil {
t.Fatalf("expected one image_generation tool, got %#v", params.Tools)
}
tool := params.Tools[0].OfImageGeneration
if tool.Model != "gpt-image-2" {
t.Fatalf("image model = %q, want gpt-image-2", tool.Model)
}
if tool.Size != "1536x1024" {
t.Fatalf("size = %q, want 1536x1024", tool.Size)
}
if tool.Quality != "medium" {
t.Fatalf("quality = %q, want medium", tool.Quality)
}
}
func TestParseCodexImageSSECompletedResponseFallback(t *testing.T) {
payload := base64.StdEncoding.EncodeToString([]byte("fake-png"))
stream := &mockCodexImageStream{
events: []responses.ResponseStreamEventUnion{{
Type: "response.completed",
Response: responses.Response{
Output: []responses.ResponseOutputItemUnion{{
Type: "image_generation_call",
Result: payload,
}},
},
}},
}
images, err := parseCodexImageSSE(stream, "png")
if err != nil {
t.Fatalf("parseCodexImageSSE: %v", err)
}
if len(images) != 1 {
t.Fatalf("images = %d, want 1", len(images))
}
if string(images[0].Data) != "fake-png" {
t.Fatalf("image data = %q, want fake-png", string(images[0].Data))
}
}

View file

@ -58,8 +58,6 @@ func NewCodexProviderWithTokenSource(
func (p *CodexProvider) Chat(
ctx context.Context, messages []Message, tools []ToolDefinition, model string, options map[string]any,
) (*LLMResponse, error) {
var opts []option.RequestOption
accountID := p.accountID
resolvedModel, fallbackReason := resolveCodexModel(model)
if fallbackReason != "" {
logger.WarnCF(
@ -72,18 +70,11 @@ func (p *CodexProvider) Chat(
},
)
}
if p.tokenSource != nil {
tok, accID, err := p.tokenSource()
if err != nil {
return nil, fmt.Errorf("refreshing token: %w", err)
}
opts = append(opts, option.WithAPIKey(tok))
if accID != "" {
accountID = accID
}
opts, accountID, err := p.requestOptions()
if err != nil {
return nil, err
}
if accountID != "" {
opts = append(opts, option.WithHeader("Chatgpt-Account-Id", accountID))
} else {
logger.WarnCF(
"provider.codex",
@ -114,7 +105,7 @@ func (p *CodexProvider) Chat(
}
}
}
err := stream.Err()
err = stream.Err()
if err != nil {
fields := map[string]any{
"requested_model": model,

View file

@ -7,17 +7,20 @@ import (
)
type (
ToolCall = protocoltypes.ToolCall
FunctionCall = protocoltypes.FunctionCall
LLMResponse = protocoltypes.LLMResponse
UsageInfo = protocoltypes.UsageInfo
Message = protocoltypes.Message
ToolDefinition = protocoltypes.ToolDefinition
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
ExtraContent = protocoltypes.ExtraContent
GoogleExtra = protocoltypes.GoogleExtra
ContentBlock = protocoltypes.ContentBlock
CacheControl = protocoltypes.CacheControl
ToolCall = protocoltypes.ToolCall
FunctionCall = protocoltypes.FunctionCall
LLMResponse = protocoltypes.LLMResponse
UsageInfo = protocoltypes.UsageInfo
Message = protocoltypes.Message
ToolDefinition = protocoltypes.ToolDefinition
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
ExtraContent = protocoltypes.ExtraContent
GoogleExtra = protocoltypes.GoogleExtra
ContentBlock = protocoltypes.ContentBlock
CacheControl = protocoltypes.CacheControl
ImageGenerationRequest = protocoltypes.ImageGenerationRequest
GeneratedImage = protocoltypes.GeneratedImage
ImageGenerationResponse = protocoltypes.ImageGenerationResponse
)
type LLMProvider interface {

View file

@ -78,6 +78,25 @@ type Attachment struct {
ContentType string `json:"content_type,omitempty"`
}
type ImageGenerationRequest struct {
Prompt string
Model string
Size string
Quality string
OutputFormat string
Count int
}
type GeneratedImage struct {
Data []byte
MimeType string
Ext string
}
type ImageGenerationResponse struct {
Images []GeneratedImage
}
type Message struct {
Role string `json:"role"`
Content string `json:"content"`

View file

@ -8,18 +8,21 @@ import (
)
type (
ToolCall = protocoltypes.ToolCall
FunctionCall = protocoltypes.FunctionCall
LLMResponse = protocoltypes.LLMResponse
UsageInfo = protocoltypes.UsageInfo
Message = protocoltypes.Message
ToolDefinition = protocoltypes.ToolDefinition
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
ExtraContent = protocoltypes.ExtraContent
GoogleExtra = protocoltypes.GoogleExtra
ContentBlock = protocoltypes.ContentBlock
CacheControl = protocoltypes.CacheControl
Attachment = protocoltypes.Attachment
ToolCall = protocoltypes.ToolCall
FunctionCall = protocoltypes.FunctionCall
LLMResponse = protocoltypes.LLMResponse
UsageInfo = protocoltypes.UsageInfo
Message = protocoltypes.Message
ToolDefinition = protocoltypes.ToolDefinition
ToolFunctionDefinition = protocoltypes.ToolFunctionDefinition
ExtraContent = protocoltypes.ExtraContent
GoogleExtra = protocoltypes.GoogleExtra
ContentBlock = protocoltypes.ContentBlock
CacheControl = protocoltypes.CacheControl
Attachment = protocoltypes.Attachment
ImageGenerationRequest = protocoltypes.ImageGenerationRequest
GeneratedImage = protocoltypes.GeneratedImage
ImageGenerationResponse = protocoltypes.ImageGenerationResponse
)
type LLMProvider interface {
@ -68,6 +71,17 @@ type NativeSearchCapable interface {
SupportsNativeSearch() bool
}
// ImageGenerationCapable is an optional interface for providers that can
// generate raster images outside the normal chat inference loop. Core tools can
// use this to reuse provider-owned auth/client behavior while keeping channel
// media delivery in the agent runtime.
type ImageGenerationCapable interface {
SupportsImageGeneration() bool
ImageGenerationProviderID() string
DefaultImageGenerationModel() string
GenerateImage(ctx context.Context, req ImageGenerationRequest) (*ImageGenerationResponse, error)
}
// FailoverReason classifies why an LLM request failed for fallback decisions.
type FailoverReason string

242
pkg/tools/image_generate.go Normal file
View file

@ -0,0 +1,242 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"github.com/google/uuid"
"github.com/sipeed/picoclaw/pkg/media"
"github.com/sipeed/picoclaw/pkg/providers"
)
const (
defaultImageGenerationSize = "1024x1024"
maxImageGenerationResults = 4
)
// ImageGenerateTool generates images through a provider adapter and returns
// generated files through the MediaStore outbound media pipeline.
type ImageGenerateTool struct {
workspace string
model string
provider providers.ImageGenerationCapable
resolver ImageGenerationProviderResolver
mediaStore media.MediaStore
}
type ImageGenerateToolOption func(*ImageGenerateTool)
type ImageGenerationProviderResolver func(model string) (providers.ImageGenerationCapable, string, error)
func WithImageGenerationProvider(provider providers.ImageGenerationCapable) ImageGenerateToolOption {
return func(t *ImageGenerateTool) {
if provider != nil {
t.provider = provider
}
}
}
func WithImageGenerationProviderResolver(resolver ImageGenerationProviderResolver) ImageGenerateToolOption {
return func(t *ImageGenerateTool) {
if resolver != nil {
t.resolver = resolver
}
}
}
func NewImageGenerateTool(
workspace string,
model string,
store media.MediaStore,
options ...ImageGenerateToolOption,
) *ImageGenerateTool {
tool := &ImageGenerateTool{
workspace: workspace,
model: model,
resolver: providers.CreateImageGenerationProviderFromModel,
mediaStore: store,
}
for _, option := range options {
option(tool)
}
return tool
}
func (t *ImageGenerateTool) SetMediaStore(store media.MediaStore) {
t.mediaStore = store
}
func (t *ImageGenerateTool) Name() string { return "image_generate" }
func (t *ImageGenerateTool) Description() string {
return `Generate an image from a prompt and send it to the current chat.
Use this when the user asks to create an image, infographic, diagram, poster, visual summary, or other generated raster artwork. The active image backend is selected from the configured image model provider prefix.`
}
func (t *ImageGenerateTool) Parameters() map[string]any {
return map[string]any{
"type": "object",
"properties": map[string]any{
"prompt": map[string]any{
"type": "string",
"description": "Image generation prompt.",
},
"size": map[string]any{
"type": "string",
"description": "Output size. Defaults to 1024x1024. Supported examples: 1024x1024, 1536x1024, 1024x1536, 2048x2048, 3840x2160.",
},
"quality": map[string]any{
"type": "string",
"enum": []string{"low", "medium", "high", "auto"},
"description": "Optional quality hint.",
},
"output_format": map[string]any{
"type": "string",
"enum": []string{"png", "jpeg", "webp"},
"description": "Output image format. Defaults to png.",
},
"count": map[string]any{
"type": "integer",
"description": "Number of images to generate, 1-4. Defaults to 1.",
},
},
"required": []string{"prompt"},
}
}
func (t *ImageGenerateTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
prompt, _ := args["prompt"].(string)
prompt = strings.TrimSpace(prompt)
if prompt == "" {
return ErrorResult("prompt is required")
}
if t.mediaStore == nil {
return ErrorResult("media store not configured")
}
if t.provider == nil && t.resolver != nil {
provider, model, err := t.resolver(t.model)
if err != nil {
return ErrorResult(fmt.Sprintf("image generation provider not configured: %v", err)).WithError(err)
}
t.provider = provider
t.model = model
}
if t.provider == nil {
return ErrorResult("image generation provider not configured")
}
req := providers.ImageGenerationRequest{
Prompt: prompt,
Model: t.model,
Size: readStringDefault(args, "size", defaultImageGenerationSize),
Quality: readStringDefault(args, "quality", ""),
OutputFormat: readStringDefault(args, "output_format", "png"),
Count: readImageCount(args["count"]),
}
if strings.TrimSpace(req.Model) == "" {
req.Model = t.provider.DefaultImageGenerationModel()
}
resp, err := t.provider.GenerateImage(ctx, req)
if err != nil {
return ErrorResult(fmt.Sprintf("image generation failed: %v", err)).WithError(err)
}
if resp == nil {
return ErrorResult("image generation returned no response")
}
images := resp.Images
if len(images) == 0 {
return ErrorResult("image generation returned no images")
}
refs := make([]string, 0, len(images))
paths := make([]string, 0, len(images))
scope := t.mediaScope(ctx)
for i, image := range images {
path, err := writeGeneratedImage(image, i)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to write generated image: %v", err)).WithError(err)
}
ref, err := t.mediaStore.Store(path, media.MediaMeta{
Filename: filepath.Base(path),
ContentType: image.MimeType,
Source: "tool:image_generate",
CleanupPolicy: media.CleanupPolicyDeleteOnCleanup,
}, scope)
if err != nil {
return ErrorResult(fmt.Sprintf("failed to register generated image: %v", err)).WithError(err)
}
refs = append(refs, ref)
paths = append(paths, path)
}
message := fmt.Sprintf("Generated %d image(s) with %s via %s.", len(refs), req.Model, t.provider.ImageGenerationProviderID())
result := MediaResult(message, refs).WithResponseHandled()
result.ArtifactTags = make([]string, 0, len(paths))
for _, path := range paths {
result.ArtifactTags = append(result.ArtifactTags, "[file:"+path+"]")
}
return result
}
func writeGeneratedImage(image providers.GeneratedImage, index int) (string, error) {
dir, err := os.MkdirTemp("", "picoclaw-image-generate-*")
if err != nil {
return "", err
}
name := fmt.Sprintf("image-%d-%s.%s", index+1, uuid.NewString(), image.Ext)
path := filepath.Join(dir, name)
if err := os.WriteFile(path, image.Data, 0o600); err != nil {
return "", err
}
return path, nil
}
func (t *ImageGenerateTool) mediaScope(ctx context.Context) string {
parts := []string{"tool:image_generate"}
if channel := ToolChannel(ctx); channel != "" {
parts = append(parts, channel)
}
if chatID := ToolChatID(ctx); chatID != "" {
parts = append(parts, chatID)
}
if sessionKey := ToolSessionKey(ctx); sessionKey != "" {
parts = append(parts, sessionKey)
}
return strings.Join(parts, ":")
}
func readStringDefault(args map[string]any, key string, fallback string) string {
value, _ := args[key].(string)
value = strings.TrimSpace(value)
if value == "" {
return fallback
}
return value
}
func readImageCount(raw any) int {
count := 1
switch v := raw.(type) {
case int:
count = v
case float64:
count = int(v)
case json.Number:
if parsed, err := v.Int64(); err == nil {
count = int(parsed)
}
}
if count < 1 {
return 1
}
if count > maxImageGenerationResults {
return maxImageGenerationResults
}
return count
}

View file

@ -0,0 +1,61 @@
package tools
import (
"context"
"testing"
"github.com/sipeed/picoclaw/pkg/media"
"github.com/sipeed/picoclaw/pkg/providers"
)
type fakeImageGenerationProvider struct {
id string
defaultModel string
request providers.ImageGenerationRequest
}
func (p *fakeImageGenerationProvider) SupportsImageGeneration() bool { return true }
func (p *fakeImageGenerationProvider) ImageGenerationProviderID() string { return p.id }
func (p *fakeImageGenerationProvider) DefaultImageGenerationModel() string { return p.defaultModel }
func (p *fakeImageGenerationProvider) GenerateImage(
_ context.Context,
req providers.ImageGenerationRequest,
) (*providers.ImageGenerationResponse, error) {
p.request = req
return &providers.ImageGenerationResponse{Images: []providers.GeneratedImage{{
Data: []byte("fake-image"),
MimeType: "image/png",
Ext: "png",
}}}, nil
}
func TestImageGenerateToolCanUseInjectedProvider(t *testing.T) {
store := media.NewFileMediaStore()
provider := &fakeImageGenerationProvider{
id: "test-provider",
defaultModel: "test-default-image-model",
}
tool := NewImageGenerateTool(
t.TempDir(),
"custom-image-model",
store,
WithImageGenerationProvider(provider),
)
result := tool.Execute(
WithToolContext(t.Context(), "telegram", "chat-1"),
map[string]any{"prompt": "make a tiny icon"},
)
if result.IsError {
t.Fatalf("Execute returned error: %s", result.ContentForLLM())
}
if provider.request.Model != "custom-image-model" {
t.Fatalf("model = %q, want custom-image-model", provider.request.Model)
}
if len(result.Media) != 1 {
t.Fatalf("media refs = %d, want 1", len(result.Media))
}
}