Merge 05e4538883 into 412705783d
This commit is contained in:
commit
4347a4f1cf
16 changed files with 893 additions and 60 deletions
|
|
@ -442,6 +442,10 @@
|
|||
"i2c": {
|
||||
"enabled": false
|
||||
},
|
||||
"image_generate": {
|
||||
"enabled": false,
|
||||
"model": "openai-codex/gpt-image-2"
|
||||
},
|
||||
"install_skill": {
|
||||
"enabled": true
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -409,6 +409,11 @@ func DefaultConfig() *Config {
|
|||
SendFile: ToolConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
ImageGenerate: ImageGenerateToolsConfig{
|
||||
ToolConfig: ToolConfig{
|
||||
Enabled: false,
|
||||
},
|
||||
},
|
||||
SendTTS: ToolConfig{
|
||||
Enabled: false,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
242
pkg/providers/oauth/codex_image_generation.go
Normal file
242
pkg/providers/oauth/codex_image_generation.go
Normal 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"
|
||||
}
|
||||
}
|
||||
111
pkg/providers/oauth/codex_image_generation_test.go
Normal file
111
pkg/providers/oauth/codex_image_generation_test.go
Normal 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))
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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"`
|
||||
|
|
|
|||
|
|
@ -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
242
pkg/tools/image_generate.go
Normal 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
|
||||
}
|
||||
61
pkg/tools/image_generate_test.go
Normal file
61
pkg/tools/image_generate_test.go
Normal 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))
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue