feat(ui): generate avatars (#80)

* feat(ui): generate avatars

Signed-off-by: mudler <mudler@localai.io>

* Show a placeholder if the image is not ready

Signed-off-by: mudler <mudler@localai.io>

* feat(avatar): generate prompt first

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: mudler <mudler@localai.io>
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/state/pool.go

@@ -2,6 +2,7 @@ package state
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -14,8 +15,11 @@ import (
 	"github.com/mudler/LocalAgent/core/agent"
 	. "github.com/mudler/LocalAgent/core/agent"
 	"github.com/mudler/LocalAgent/core/sse"
+	"github.com/mudler/LocalAgent/pkg/llm"
 	"github.com/mudler/LocalAgent/pkg/localrag"
 	"github.com/mudler/LocalAgent/pkg/utils"
+	"github.com/sashabaranov/go-openai"
+	"github.com/sashabaranov/go-openai/jsonschema"
 
 	"github.com/mudler/LocalAgent/pkg/xlog"
 )
@@ -153,9 +157,88 @@ func (a *AgentPool) CreateAgent(name string, agentConfig *AgentConfig) error {
 		return err
 	}
 
+	go func() {
+		// Create the agent avatar
+		if err := a.createAgentAvatar(agentConfig); err != nil {
+			xlog.Error("Failed to create agent avatar", "error", err)
+		}
+	}()
+
 	return a.startAgentWithConfig(name, agentConfig)
 }
 
+func (a *AgentPool) createAgentAvatar(agent *AgentConfig) error {
+	client := llm.NewClient(a.apiKey, a.apiURL+"/v1", "10m")
+
+	if a.imageModel == "" {
+		return fmt.Errorf("image model not set")
+	}
+
+	if a.defaultModel == "" {
+		return fmt.Errorf("default model not set")
+	}
+
+	var results struct {
+		ImagePrompt string `json:"image_prompt"`
+	}
+
+	err := llm.GenerateTypedJSON(
+		context.Background(),
+		llm.NewClient(a.apiKey, a.apiURL, "10m"),
+		"Generate a prompt that I can use to create a random avatar for the bot '"+agent.Name+"', the description of the bot is: "+agent.Description,
+		a.defaultModel,
+		jsonschema.Definition{
+			Type: jsonschema.Object,
+			Properties: map[string]jsonschema.Definition{
+				"image_prompt": {
+					Type:        jsonschema.String,
+					Description: "The prompt to generate the image",
+				},
+			},
+			Required: []string{"image_prompt"},
+		}, &results)
+	if err != nil {
+		return fmt.Errorf("failed to generate image prompt: %w", err)
+	}
+
+	if results.ImagePrompt == "" {
+		xlog.Error("Failed to generate image prompt")
+		return fmt.Errorf("failed to generate image prompt")
+	}
+
+	req := openai.ImageRequest{
+		Prompt:         results.ImagePrompt,
+		Model:          a.imageModel,
+		Size:           openai.CreateImageSize256x256,
+		ResponseFormat: openai.CreateImageResponseFormatB64JSON,
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	resp, err := client.CreateImage(ctx, req)
+	if err != nil {
+		return fmt.Errorf("failed to generate image: %w", err)
+	}
+
+	if len(resp.Data) == 0 {
+		return fmt.Errorf("failed to generate image")
+	}
+
+	imageJson := resp.Data[0].B64JSON
+
+	os.MkdirAll(filepath.Join(a.pooldir, "avatars"), 0755)
+
+	// Save the image to the agent directory
+	imagePath := filepath.Join(a.pooldir, "avatars", fmt.Sprintf("%s.png", agent.Name))
+	imageData, err := base64.StdEncoding.DecodeString(imageJson)
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(imagePath, imageData, 0644)
+}
+
 func (a *AgentPool) List() []string {
 	a.Lock()
 	defer a.Unlock()
@@ -474,6 +557,10 @@ func (a *AgentPool) Remove(name string) error {
 	a.stop(name)
 	delete(a.agents, name)
 	delete(a.pool, name)
+
+	// remove avatar
+	os.Remove(filepath.Join(a.pooldir, "avatars", fmt.Sprintf("%s.png", name)))
+
 	if err := a.save(); err != nil {
 		return err
 	}

main.go

@@ -77,6 +77,7 @@ func main() {
 		webui.WithLLMAPIUrl(apiURL),
 		webui.WithLLMAPIKey(apiKey),
 		webui.WithLLMModel(testModel),
+		webui.WithStateDir(stateDir),
 	)
 
 	// Start the agents

webui/options.go

@@ -9,6 +9,7 @@ type Config struct {
 	LLMAPIURL        string
 	LLMAPIKey        string
 	LLMModel         string
+	StateDir         string
 }
 
 type Option func(*Config)
@@ -19,6 +20,12 @@ func WithDefaultChunkSize(size int) Option {
 	}
 }
 
+func WithStateDir(dir string) Option {
+	return func(c *Config) {
+		c.StateDir = dir
+	}
+}
+
 func WithLLMModel(model string) Option {
 	return func(c *Config) {
 		c.LLMModel = model

webui/routes.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"math/rand"
 	"net/http"
+	"path/filepath"
 
 	"github.com/dave-gray101/v2keyauth"
 	fiber "github.com/gofiber/fiber/v2"
@@ -26,6 +27,13 @@ var embeddedFiles embed.FS
 
 func (app *App) registerRoutes(pool *state.AgentPool, webapp *fiber.App) {
 
+	// Static avatars in a.pooldir/avatars
+	webapp.Use("/avatars", filesystem.New(filesystem.Config{
+		Root: http.Dir(filepath.Join(app.config.StateDir, "avatars")),
+		//	PathPrefix: "avatars",
+		Browse: true,
+	}))
+
 	webapp.Use("/public", filesystem.New(filesystem.Config{
 		Root:       http.FS(embeddedFiles),
 		PathPrefix: "public",
@@ -145,6 +153,7 @@ func (app *App) registerRoutes(pool *state.AgentPool, webapp *fiber.App) {
 
 	webapp.Post("/settings/import", app.ImportAgent(pool))
 	webapp.Get("/settings/export/:name", app.ExportAgent(pool))
+
 }
 
 var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")

webui/views/agents.html

@@ -5,6 +5,45 @@
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>Agent List</title>
     {{template "views/partials/header"}}
+    <style>
+        .avatar-placeholder {
+            width: 96px;
+            height: 96px;
+            border-radius: 50%;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            background: linear-gradient(135deg, #2a2a2a, #3a3a3a);
+            color: var(--primary);
+            font-size: 1.5rem;
+            margin-bottom: 1rem;
+            border: 2px solid var(--primary);
+            box-shadow: var(--neon-glow);
+            position: relative;
+            overflow: hidden;
+        }
+        
+        .avatar-placeholder::after {
+            content: "";
+            position: absolute;
+            width: 100%;
+            height: 4px;
+            background: var(--primary);
+            bottom: 0;
+            left: 0;
+            animation: loading-progress 2s infinite linear;
+        }
+        
+        @keyframes loading-progress {
+            0% { width: 0; }
+            50% { width: 100%; }
+            100% { width: 0; }
+        }
+        
+        .placeholder-text {
+            z-index: 1;
+        }
+    </style>
 </head>
 <body>
     {{template "views/partials/menu"}}
@@ -65,8 +104,15 @@
             {{ range .Agents }}
             <div hx-ext="sse" data-agent-name="{{.}}" class="card">
                 <div class="flex flex-col items-center text-center p-4">
-                    <!--<img src="/public/agent_{{.}}.png" alt="{{.}}" class="w-24 h-24 rounded-full mb-4" 
+                    <div class="avatar-container mb-4">
-                         style="border: 2px solid var(--primary); box-shadow: var(--neon-glow);">-->
+                        <img src="/avatars/{{.}}.png" alt="{{.}}" class="w-24 h-24 rounded-full" 
+                             style="border: 2px solid var(--primary); box-shadow: var(--neon-glow); display: none;"
+                             onload="this.style.display = 'block'; this.nextElementSibling.style.display = 'none';"
+                             onerror="this.style.display = 'none'; this.nextElementSibling.style.display = 'flex';">
+                        <div class="avatar-placeholder">
+                            <span class="placeholder-text"><i class="fas fa-sync fa-spin"></i></span>
+                        </div>
+                    </div>
                     <h2>{{.}}</h2>
                     <div class="mb-4 flex items-center justify-center">
                         <span class="badge {{ if eq (index $status .) true }}badge-primary{{ else }}badge-secondary{{ end }} mr-2">
@@ -117,6 +163,22 @@
 
     <script>
         document.addEventListener('DOMContentLoaded', function() {
+            // Image loading handler
+            document.querySelectorAll('.avatar-container img').forEach(img => {
+                // Check if image is already cached
+                if (img.complete) {
+                    if (img.naturalHeight === 0) {
+                        // Image failed to load
+                        img.style.display = 'none';
+                        img.nextElementSibling.style.display = 'flex';
+                    } else {
+                        // Image loaded successfully
+                        img.style.display = 'block';
+                        img.nextElementSibling.style.display = 'none';
+                    }
+                }
+                // onload and onerror handlers are already in the HTML
+            });
     
             const importSection = document.getElementById('import-section');
             const toggleImport = document.getElementById('toggle-import');