feat: RAG, memory, plugins, lessons, file editor, split panes, Markdown rendering, PWA + UI overhaul
All checks were successful
Beta Release / beta (push) Successful in 5m9s
All checks were successful
Beta Release / beta (push) Successful in 5m9s
Major additions: - RAG pipeline (indexing, chunking, search) with sidebar upload button - Memory system with CRUD API - Plugins and lessons modules - MCP discovery and MCP server - Advanced skills (auto-create, conditional, improver) - Agent browser/image support, delegate, sessions - File editor with CodeMirror in split panes - Markdown rendering via react-markdown + KaTeX + highlight.js - Raw markdown toggle - PWA manifest + service worker - Extension UI redesign with new design tokens and studio-style chat - Pipeline API for chat streaming - Mobile responsive layout 💘 Generated with Crush Assisted-by: GLM-5.1 via Crush <crush@charm.land>
This commit is contained in:
200
internal/agent/image.go
Normal file
200
internal/agent/image.go
Normal file
@@ -0,0 +1,200 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/muyue/muyue/internal/config"
|
||||
)
|
||||
|
||||
type ImageGenerationTool struct {
|
||||
apiKey string
|
||||
baseURL string
|
||||
model string
|
||||
saveDir string
|
||||
}
|
||||
|
||||
func NewImageGenerationTool(cfg *config.MuyueConfig) (*ImageGenerationTool, error) {
|
||||
configDir, err := config.ConfigDir()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
saveDir := filepath.Join(configDir, "images")
|
||||
if err := os.MkdirAll(saveDir, 0755); err != nil {
|
||||
return nil, fmt.Errorf("creating images dir: %w", err)
|
||||
}
|
||||
|
||||
var apiKey, baseURL, model string
|
||||
for _, p := range cfg.AI.Providers {
|
||||
if p.Active {
|
||||
apiKey = p.APIKey
|
||||
baseURL = p.BaseURL
|
||||
model = p.Model
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if baseURL == "" {
|
||||
baseURL = "https://api.openai.com/v1"
|
||||
}
|
||||
|
||||
return &ImageGenerationTool{
|
||||
apiKey: apiKey,
|
||||
baseURL: strings.TrimRight(baseURL, "/"),
|
||||
model: model,
|
||||
saveDir: saveDir,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (t *ImageGenerationTool) Name() string {
|
||||
return "generate_image"
|
||||
}
|
||||
|
||||
func (t *ImageGenerationTool) Description() string {
|
||||
return "Generate an image from a text prompt using DALL-E or compatible API. Returns a local URL to the generated image."
|
||||
}
|
||||
|
||||
func (t *ImageGenerationTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"prompt": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Description of the image to generate",
|
||||
},
|
||||
"size": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Image size: 1024x1024, 1024x1792, or 1792x1024",
|
||||
"default": "1024x1024",
|
||||
},
|
||||
"style": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Style: vivid or natural",
|
||||
"default": "vivid",
|
||||
},
|
||||
},
|
||||
"required": []string{"prompt"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *ImageGenerationTool) Execute(args map[string]interface{}) (string, error) {
|
||||
prompt, _ := args["prompt"].(string)
|
||||
if prompt == "" {
|
||||
return "", fmt.Errorf("prompt is required")
|
||||
}
|
||||
|
||||
size, _ := args["size"].(string)
|
||||
if size == "" {
|
||||
size = "1024x1024"
|
||||
}
|
||||
style, _ := args["style"].(string)
|
||||
if style == "" {
|
||||
style = "vivid"
|
||||
}
|
||||
|
||||
reqBody := map[string]interface{}{
|
||||
"model": "dall-e-3",
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
"style": style,
|
||||
"n": 1,
|
||||
}
|
||||
|
||||
bodyBytes, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("marshal request: %w", err)
|
||||
}
|
||||
|
||||
url := t.baseURL + "/images/generations"
|
||||
req, err := http.NewRequest("POST", url, bytes.NewReader(bodyBytes))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if t.apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+t.apiKey)
|
||||
}
|
||||
|
||||
client := &http.Client{Timeout: 120 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("API error (%d): %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
|
||||
var genResp struct {
|
||||
Data []struct {
|
||||
URL string `json:"url"`
|
||||
B64JSON string `json:"b64_json"`
|
||||
RevisedPrompt string `json:"revised_prompt"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(respBody, &genResp); err != nil {
|
||||
return "", fmt.Errorf("parse response: %w", err)
|
||||
}
|
||||
|
||||
if len(genResp.Data) == 0 {
|
||||
return "", fmt.Errorf("no image returned")
|
||||
}
|
||||
|
||||
imgData := genResp.Data[0]
|
||||
filename := fmt.Sprintf("img-%d.png", time.Now().UnixNano())
|
||||
localPath := filepath.Join(t.saveDir, filename)
|
||||
|
||||
if imgData.B64JSON != "" {
|
||||
return "", fmt.Errorf("base64 response not yet supported")
|
||||
}
|
||||
|
||||
if imgData.URL != "" {
|
||||
if err := t.downloadImage(imgData.URL, localPath); err != nil {
|
||||
return "", fmt.Errorf("download image: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
result := map[string]interface{}{
|
||||
"url": "/api/images/" + filename,
|
||||
"revised_prompt": imgData.RevisedPrompt,
|
||||
"size": size,
|
||||
}
|
||||
resultJSON, _ := json.Marshal(result)
|
||||
return string(resultJSON), nil
|
||||
}
|
||||
|
||||
func (t *ImageGenerationTool) downloadImage(url, localPath string) error {
|
||||
resp, err := http.Get(url)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("download failed: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
f, err := os.Create(localPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
_, err = io.Copy(f, resp.Body)
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user