Merge pull request #816 from trheyi/main

Refactor Neo API and Enhance Assistant Capabilities (Dev)

neo/api.go

@@ -15,6 +15,7 @@ import (
 	"github.com/yaoapp/gou/connector"
 	"github.com/yaoapp/gou/process"
 	"github.com/yaoapp/yao/helper"
+	chatctx "github.com/yaoapp/yao/neo/context"
 	"github.com/yaoapp/yao/neo/message"
 	"github.com/yaoapp/yao/neo/store"
 )
@@ -173,7 +174,7 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 	}
 
 	// Set the context
-	ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
+	ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
 	defer cancel()
 
 	// Upload the file
@@ -214,7 +215,7 @@ func (neo *DSL) handleChat(c *gin.Context) {
 	}
 
 	// Set the context with validated chat_id
-	ctx, cancel := NewContextWithCancel(sid, chatID, c.Query("context"))
+	ctx, cancel := chatctx.NewWithCancel(sid, chatID, c.Query("context"))
 	defer cancel()
 
 	neo.Answer(ctx, content, c)
@@ -297,7 +298,7 @@ func (neo *DSL) handleDownload(c *gin.Context) {
 	}
 
 	// Set the context
-	ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
+	ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
 	defer cancel()
 
 	// Download the file
@@ -537,7 +538,7 @@ func (neo *DSL) handleChatUpdate(c *gin.Context) {
 
 	// If content is not empty, Generate the chat title
 	if body.Content != "" {
-		ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
+		ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
 		defer cancel()
 
 		title, err := neo.GenerateChatTitle(ctx, body.Content, c, true)
@@ -729,7 +730,7 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) {
 		return
 	}
 
-	ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
+	ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
 	defer cancel()
 
 	// Use silent mode for regular HTTP requests, streaming for SSE
@@ -780,7 +781,7 @@ func (neo *DSL) handleGeneratePrompts(c *gin.Context) {
 		return
 	}
 
-	ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
+	ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
 	defer cancel()
 
 	// Use silent mode for regular HTTP requests, streaming for SSE
@@ -831,7 +832,7 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) {
 		return
 	}
 
-	ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
+	ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
 	defer cancel()
 
 	// Use silent mode for regular HTTP requests, streaming for SSE

neo/assistant/api.go

@@ -2,16 +2,14 @@ package assistant
 
 import (
 	"context"
-	"crypto/sha256"
 	"encoding/base64"
 	"fmt"
-	"io"
-	"mime/multipart"
-	"path/filepath"
 	"strings"
-	"time"
 
+	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/gou/fs"
+	chatctx "github.com/yaoapp/yao/neo/context"
+	"github.com/yaoapp/yao/neo/message"
 	chatMessage "github.com/yaoapp/yao/neo/message"
 )
 
@@ -45,24 +43,238 @@ func GetByConnector(connector string, name string) (*Assistant, error) {
 	return assistant, nil
 }
 
-// AllowedFileTypes the allowed file types
+// Execute implements the execute functionality
-var AllowedFileTypes = map[string]string{
+func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error {
-	"application/json":   "json",
+	messages, err := ast.withHistory(ctx, input)
-	"application/pdf":    "pdf",
+	if err != nil {
-	"application/msword": "doc",
+		return err
-	"application/vnd.openxmlformats-officedocument.wordprocessingml.document":   "docx",
-	"application/vnd.oasis.opendocument.text":                                   "odt",
-	"application/vnd.ms-excel":                                                  "xls",
-	"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":         "xlsx",
-	"application/vnd.ms-powerpoint":                                             "ppt",
-	"application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
 	}
 
-// MaxSize 20M max file size
+	options = ast.withOptions(options)
-var MaxSize int64 = 20 * 1024 * 1024
+
+	// Run init hook
+	res, err := ast.HookInit(c, ctx, messages, options)
+	if err != nil {
+		return err
+	}
+
+	// Switch to the new assistant if necessary
+	if res.AssistantID != ctx.AssistantID {
+		newAst, err := Get(res.AssistantID)
+		if err != nil {
+			return err
+		}
+		*ast = *newAst
+	}
+
+	// Handle next action
+	if res.Next != nil {
+		switch res.Next.Action {
+		case "exit":
+			return nil
+			// Add other actions here if needed
+		}
+	}
+
+	// Update options if provided
+	if res.Options != nil {
+		options = res.Options
+	}
+
+	// messages
+	if res.Input != nil {
+		messages = res.Input
+	}
+
+	// Only proceed with chat stream if no specific next action was handled
+	return ast.handleChatStream(c, ctx, messages, options)
+}
+
+// handleChatStream manages the streaming chat interaction with the AI
+func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{}) error {
+	clientBreak := make(chan bool, 1)
+	done := make(chan bool, 1)
+	content := []byte{}
+
+	// Chat with AI in background
+	go func() {
+		err := ast.streamChat(c, ctx, messages, options, clientBreak, done, &content)
+		if err != nil {
+			chatMessage.New().Error(err).Done().Write(c.Writer)
+		}
+
+		ast.saveChatHistory(ctx, messages, content)
+		done <- true
+	}()
+
+	// Wait for completion or client disconnect
+	select {
+	case <-done:
+		return nil
+	case <-c.Writer.CloseNotify():
+		clientBreak <- true
+		return nil
+	}
+}
+
+// streamChat handles the streaming chat interaction
+func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{},
+	clientBreak chan bool, done chan bool, content *[]byte) error {
+
+	return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
+		select {
+		case <-clientBreak:
+			return 0 // break
+
+		default:
+			msg := chatMessage.NewOpenAI(data)
+			if msg == nil {
+				return 1 // continue
+			}
+
+			// Handle error
+			if msg.Type == "error" {
+				value := msg.String()
+				res, hookErr := ast.HookFail(c, ctx, messages, string(*content), fmt.Errorf("%s", value))
+				if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
+					value = res.Output
+					if res.Error != "" {
+						value = res.Error
+					}
+				}
+				chatMessage.New().Error(value).Done().Write(c.Writer)
+				return 0 // break
+			}
+
+			// Append content and send message
+			*content = msg.Append(*content)
+			value := msg.String()
+			if value != "" {
+				// Handle stream
+				res, err := ast.HookStream(c, ctx, messages, string(*content))
+				if err == nil && res != nil {
+					if res.Output != "" {
+						value = res.Output
+					}
+					if res.Next != nil && res.Next.Action == "exit" {
+						done <- true
+						return 0 // break
+					}
+					if res.Silent {
+						return 1 // continue
+					}
+				}
+
+				chatMessage.New().
+					Map(map[string]interface{}{
+						"text": value,
+						"done": msg.IsDone,
+					}).
+					Write(c.Writer)
+			}
+
+			// Complete the stream
+			if msg.IsDone {
+				// if value == "" {
+				// 	msg.Write(c.Writer)
+				// }
+
+				// Call HookDone
+				res, hookErr := ast.HookDone(c, ctx, messages, string(*content))
+				if hookErr == nil && res != nil {
+					if res.Output != "" {
+						chatMessage.New().
+							Map(map[string]interface{}{
+								"text": res.Output,
+								"done": true,
+							}).
+							Write(c.Writer)
+					}
+					if res.Next != nil && res.Next.Action == "exit" {
+						done <- true
+						return 0 // break
+					}
+				} else if value != "" {
+					chatMessage.New().
+						Map(map[string]interface{}{
+							"text": value,
+							"done": true,
+						}).
+						Write(c.Writer)
+				}
+
+				done <- true
+				return 0 // break
+			}
+
+			return 1 // continue
+		}
+	})
+}
+
+// saveChatHistory saves the chat history if storage is available
+func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content []byte) {
+	if len(content) > 0 && ctx.Sid != "" && len(messages) > 0 {
+		storage.SaveHistory(
+			ctx.Sid,
+			[]map[string]interface{}{
+				{"role": "user", "content": messages[len(messages)-1].Content(), "name": ctx.Sid},
+				{"role": "assistant", "content": string(content), "name": ctx.Sid},
+			},
+			ctx.ChatID,
+			nil,
+		)
+	}
+}
+
+func (ast *Assistant) withOptions(options map[string]interface{}) map[string]interface{} {
+	if options == nil {
+		options = map[string]interface{}{}
+	}
+
+	if ast.Options != nil {
+		for key, value := range ast.Options {
+			options[key] = value
+		}
+	}
+	return options
+}
+
+func (ast *Assistant) withPrompts(messages []message.Message) []message.Message {
+	if ast.Prompts != nil {
+		for _, prompt := range ast.Prompts {
+			name := ast.Name
+			if prompt.Name != "" {
+				name = prompt.Name
+			}
+			messages = append(messages, *message.New().Map(map[string]interface{}{"role": prompt.Role, "content": prompt.Content, "name": name}))
+		}
+	}
+	return messages
+}
+
+func (ast *Assistant) withHistory(ctx chatctx.Context, input string) ([]message.Message, error) {
+	messages := []message.Message{}
+	messages = ast.withPrompts(messages)
+	if storage != nil {
+		history, err := storage.GetHistory(ctx.Sid, ctx.ChatID)
+		if err != nil {
+			return nil, err
+		}
+
+		// Add history messages
+		for _, h := range history {
+			messages = append(messages, *message.New().Map(h))
+		}
+	}
+
+	// Add user message
+	messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": input, "name": ctx.Sid}))
+	return messages, nil
+}
 
 // Chat implements the chat functionality
-func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error {
+func (ast *Assistant) Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error {
 	if ast.openai == nil {
 		return fmt.Errorf("openai is not initialized")
 	}
@@ -80,13 +292,12 @@ func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{
 	return nil
 }
 
-func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
+func (ast *Assistant) requestMessages(ctx context.Context, messages []message.Message) ([]map[string]interface{}, error) {
 	newMessages := []map[string]interface{}{}
-
 	// With Prompts
 	if ast.Prompts != nil {
 		for _, prompt := range ast.Prompts {
-			message := map[string]interface{}{
+			msg := map[string]interface{}{
 				"role":    prompt.Role,
 				"content": prompt.Content,
 			}
@@ -96,20 +307,20 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string
 				name = prompt.Name
 			}
 
-			message["name"] = name
+			msg["name"] = name
-			newMessages = append(newMessages, message)
+			newMessages = append(newMessages, msg)
 		}
 	}
 
 	length := len(messages)
 	for index, message := range messages {
-		role, ok := message["role"].(string)
+		role := message.Role
-		if !ok {
+		if role == "" {
 			return nil, fmt.Errorf("role must be string")
 		}
 
-		content, ok := message["content"].(string)
+		content := message.Text
-		if !ok {
+		if content == "" {
 			return nil, fmt.Errorf("content must be string")
 		}
 
@@ -118,7 +329,7 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string
 			"content": content,
 		}
 
-		if name, ok := message["name"].(string); ok {
+		if name := message.Name; name != "" {
 			newMessage["name"] = name
 		}
 
@@ -175,94 +386,6 @@ func (ast *Assistant) withAttachments(ctx context.Context, msg *chatMessage.Mess
 	return contents, nil
 }
 
-// Upload implements file upload functionality
-func (ast *Assistant) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error) {
-	// check file size
-	if file.Size > MaxSize {
-		return nil, fmt.Errorf("file size %d exceeds the maximum size of %d", file.Size, MaxSize)
-	}
-
-	contentType := file.Header.Get("Content-Type")
-	if !ast.allowed(contentType) {
-		return nil, fmt.Errorf("file type %s not allowed", contentType)
-	}
-
-	data, err := fs.Get("data")
-	if err != nil {
-		return nil, err
-	}
-
-	ext := filepath.Ext(file.Filename)
-	id, err := ast.id(file.Filename, ext)
-	if err != nil {
-		return nil, err
-	}
-
-	filename := id
-	_, err = data.Write(filename, reader, 0644)
-	if err != nil {
-		return nil, err
-	}
-
-	return &File{
-		ID:          filename,
-		Filename:    filename,
-		ContentType: contentType,
-		Bytes:       int(file.Size),
-		CreatedAt:   int(time.Now().Unix()),
-	}, nil
-}
-
-func (ast *Assistant) allowed(contentType string) bool {
-	if _, ok := AllowedFileTypes[contentType]; ok {
-		return true
-	}
-	if strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "image/") ||
-		strings.HasPrefix(contentType, "audio/") || strings.HasPrefix(contentType, "video/") {
-		return true
-	}
-	return false
-}
-
-func (ast *Assistant) id(temp string, ext string) (string, error) {
-	date := time.Now().Format("20060102")
-	hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8]
-	return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil
-}
-
-// Download implements file download functionality
-func (ast *Assistant) Download(ctx context.Context, fileID string) (*FileResponse, error) {
-	data, err := fs.Get("data")
-	if err != nil {
-		return nil, fmt.Errorf("get filesystem error: %s", err.Error())
-	}
-
-	exists, err := data.Exists(fileID)
-	if err != nil {
-		return nil, fmt.Errorf("check file error: %s", err.Error())
-	}
-	if !exists {
-		return nil, fmt.Errorf("file %s not found", fileID)
-	}
-
-	reader, err := data.ReadCloser(fileID)
-	if err != nil {
-		return nil, err
-	}
-
-	ext := filepath.Ext(fileID)
-	contentType := "application/octet-stream"
-	if v, err := data.MimeType(fileID); err == nil {
-		contentType = v
-	}
-
-	return &FileResponse{
-		Reader:      reader,
-		ContentType: contentType,
-		Extension:   ext,
-	}, nil
-}
-
 // ReadBase64 implements base64 file reading functionality
 func (ast *Assistant) ReadBase64(ctx context.Context, fileID string) (string, error) {
 	data, err := fs.Get("data")

neo/assistant/attachment.go

@@ -0,0 +1,322 @@
+package assistant
+
+import (
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"fmt"
+	"io"
+	"mime/multipart"
+	"path/filepath"
+	"strings"
+	"time"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/yaoapp/gou/fs"
+	"github.com/yaoapp/gou/rag/driver"
+)
+
+// AllowedFileTypes the allowed file types
+var AllowedFileTypes = map[string]string{
+	"application/json":   "json",
+	"application/pdf":    "pdf",
+	"application/msword": "doc",
+	"application/vnd.openxmlformats-officedocument.wordprocessingml.document":   "docx",
+	"application/vnd.oasis.opendocument.text":                                   "odt",
+	"application/vnd.ms-excel":                                                  "xls",
+	"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":         "xlsx",
+	"application/vnd.ms-powerpoint":                                             "ppt",
+	"application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
+}
+
+// MaxSize 20M max file size
+var MaxSize int64 = 20 * 1024 * 1024
+
+// Upload implements file upload functionality
+func (ast *Assistant) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error) {
+	// check file size
+	if file.Size > MaxSize {
+		return nil, fmt.Errorf("file size %d exceeds the maximum size of %d", file.Size, MaxSize)
+	}
+
+	contentType := file.Header.Get("Content-Type")
+	if !ast.allowed(contentType) {
+		return nil, fmt.Errorf("file type %s not allowed", contentType)
+	}
+
+	// Get chat ID and session ID from options
+	chatID := ""
+	sid := ""
+	if v, ok := option["chat_id"].(string); ok {
+		chatID = v
+	}
+	if v, ok := option["sid"].(string); ok {
+		sid = v
+	}
+
+	// Generate file ID with namespace
+	fileID, err := ast.generateFileID(file.Filename, sid, chatID)
+	if err != nil {
+		return nil, err
+	}
+
+	// Upload file to storage
+	data, err := fs.Get("data")
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = data.Write(fileID, reader, 0644)
+	if err != nil {
+		return nil, err
+	}
+
+	// Create file response
+	fileResp := &File{
+		ID:          fileID,
+		Filename:    fileID,
+		ContentType: contentType,
+		Bytes:       int(file.Size),
+		CreatedAt:   int(time.Now().Unix()),
+	}
+
+	// Handle RAG if available
+	if err := ast.handleRAG(ctx, fileResp, reader, option); err != nil {
+		return nil, fmt.Errorf("RAG handling error: %s", err.Error())
+	}
+
+	// Handle Vision if available
+	if err := ast.handleVision(ctx, fileResp, option); err != nil {
+		return nil, fmt.Errorf("Vision handling error: %s", err.Error())
+	}
+
+	return fileResp, nil
+}
+
+// generateFileID generates a file ID with proper namespace
+func (ast *Assistant) generateFileID(filename string, sid string, chatID string) (string, error) {
+	ext := filepath.Ext(filename)
+	hash := fmt.Sprintf("%x", sha256.Sum256([]byte(filename)))[:8]
+	date := time.Now().Format("20060102")
+
+	// Build namespace
+	namespace := fmt.Sprintf("__assistants/%s", ast.ID)
+	if sid != "" {
+		namespace = fmt.Sprintf("%s/%s", namespace, sid)
+		if chatID != "" {
+			namespace = fmt.Sprintf("%s/%s", namespace, chatID)
+		}
+	}
+
+	return fmt.Sprintf("%s/%s/%s%s", namespace, date, hash, ext), nil
+}
+
+// handleRAG handles the file with RAG if available
+func (ast *Assistant) handleRAG(ctx context.Context, file *File, reader io.Reader, option map[string]interface{}) error {
+	if rag == nil {
+		return nil
+	}
+
+	// Check if RAG processing is enabled
+	if option, ok := option["rag"].(bool); !ok || !option {
+		return nil
+	}
+
+	// Only handle text-based files
+	if !strings.HasPrefix(file.ContentType, "text/") {
+		return nil
+	}
+
+	// Reset reader to beginning
+	if seeker, ok := reader.(io.Seeker); ok {
+		if _, err := seeker.Seek(0, io.SeekStart); err != nil {
+			return err
+		}
+	}
+
+	// Extract sid and chat_id from file path
+	parts := strings.Split(file.ID, "/")
+	indexName := fmt.Sprintf("%s%s", rag.Setting.IndexPrefix, ast.ID) // Default: prefix-assistant
+
+	if len(parts) >= 4 { // Has sid
+		sid := parts[2]
+		indexName = fmt.Sprintf("%s%s-%s", rag.Setting.IndexPrefix, ast.ID, sid) // prefix-assistant-user
+
+		if len(parts) >= 5 { // Has chat_id
+			chatID := parts[3]
+			indexName = fmt.Sprintf("%s%s-%s-%s", rag.Setting.IndexPrefix, ast.ID, sid, chatID) // prefix-assistant-user-chat
+		}
+	}
+
+	// Check if index exists
+	exists, err := rag.Engine.HasIndex(ctx, indexName)
+	if err != nil {
+		return fmt.Errorf("check index error: %s", err.Error())
+	}
+
+	// Create index if not exists
+	if !exists {
+		err = rag.Engine.CreateIndex(ctx, driver.IndexConfig{Name: indexName})
+		if err != nil {
+			return fmt.Errorf("create index error: %s", err.Error())
+		}
+	}
+
+	// Reset reader again after checking index
+	if seeker, ok := reader.(io.Seeker); ok {
+		if _, err := seeker.Seek(0, io.SeekStart); err != nil {
+			return err
+		}
+	}
+
+	// Upload and index the file
+	result, err := rag.Uploader.Upload(ctx, reader, driver.FileUploadOptions{
+		Async:        false,
+		ChunkSize:    1024, // Default chunk size
+		ChunkOverlap: 256,  // Default overlap
+		IndexName:    indexName,
+	})
+
+	if err != nil {
+		return fmt.Errorf("upload error: %s", err.Error())
+	}
+
+	if len(result.Documents) == 0 {
+		return fmt.Errorf("no documents indexed")
+	}
+
+	// Store the document IDs
+	docIDs := make([]string, len(result.Documents))
+	for i, doc := range result.Documents {
+		docIDs[i] = doc.DocID
+	}
+	file.DocIDs = docIDs
+
+	return nil
+}
+
+// handleVision handles the file with Vision if available
+func (ast *Assistant) handleVision(ctx context.Context, file *File, option map[string]interface{}) error {
+	if vision == nil {
+		return nil
+	}
+
+	// Check if Vision processing is enabled
+	if option, ok := option["vision"].(bool); !ok || !option {
+		return nil
+	}
+
+	// Check if file is an image
+	if !strings.HasPrefix(file.ContentType, "image/") {
+		return nil
+	}
+
+	// Get model from options
+	model := ""
+	if v, ok := option["model"].(string); ok {
+		model = v
+	}
+
+	// Reset reader for vision service
+	data, err := fs.Get("data")
+	if err != nil {
+		return fmt.Errorf("get filesystem error: %s", err.Error())
+	}
+
+	exists, err := data.Exists(file.ID)
+	if err != nil {
+		return fmt.Errorf("check file error: %s", err.Error())
+	}
+	if !exists {
+		return fmt.Errorf("file %s not found", file.ID)
+	}
+
+	// Read file content into memory
+	imgData, err := data.ReadFile(file.ID)
+	if err != nil {
+		return fmt.Errorf("read file error: %s", err.Error())
+	}
+
+	if VisionCapableModels[model] {
+		// For vision-capable models, upload to vision service to get URL
+		resp, err := vision.Upload(ctx, file.Filename, bytes.NewReader(imgData), file.ContentType)
+		if err != nil {
+			return fmt.Errorf("vision upload error: %s", err.Error())
+		}
+		file.URL = resp.URL // Store the URL for vision-capable models to use
+	} else {
+		// For non-vision models, get image description
+		prompt := "Describe this image in detail."
+		if v, ok := option["vision_prompt"].(string); ok {
+			prompt = v
+		}
+
+		// Upload to vision service first Compress image
+		resp, err := vision.Upload(ctx, file.Filename, bytes.NewReader(imgData), file.ContentType)
+		if err != nil {
+			return fmt.Errorf("vision upload error: %s", err.Error())
+		}
+
+		// Analyze using base64 data
+		result, err := vision.Analyze(ctx, resp.FileID, prompt)
+		if err != nil {
+			return fmt.Errorf("vision analyze error: %s", err.Error())
+		}
+
+		// Extract description text from response
+		if desc, ok := result.Description["text"].(string); ok {
+			file.Description = desc
+		} else {
+			// Convert the entire description to JSON string as fallback
+			bytes, err := jsoniter.Marshal(result.Description)
+			if err == nil {
+				file.Description = string(bytes)
+			}
+		}
+	}
+	return nil
+}
+
+// Download implements file download functionality
+func (ast *Assistant) Download(ctx context.Context, fileID string) (*FileResponse, error) {
+	data, err := fs.Get("data")
+	if err != nil {
+		return nil, fmt.Errorf("get filesystem error: %s", err.Error())
+	}
+
+	exists, err := data.Exists(fileID)
+	if err != nil {
+		return nil, fmt.Errorf("check file error: %s", err.Error())
+	}
+	if !exists {
+		return nil, fmt.Errorf("file %s not found", fileID)
+	}
+
+	reader, err := data.ReadCloser(fileID)
+	if err != nil {
+		return nil, err
+	}
+
+	ext := filepath.Ext(fileID)
+	contentType := "application/octet-stream"
+	if v, err := data.MimeType(fileID); err == nil {
+		contentType = v
+	}
+
+	return &FileResponse{
+		Reader:      reader,
+		ContentType: contentType,
+		Extension:   ext,
+	}, nil
+}
+
+func (ast *Assistant) allowed(contentType string) bool {
+	if _, ok := AllowedFileTypes[contentType]; ok {
+		return true
+	}
+	if strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "image/") ||
+		strings.HasPrefix(contentType, "audio/") || strings.HasPrefix(contentType, "video/") {
+		return true
+	}
+	return false
+}

neo/assistant/attachment_test.go

@@ -0,0 +1,363 @@
+package assistant
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"io"
+	"mime/multipart"
+	"os"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/yaoapp/gou/fs"
+	gourag "github.com/yaoapp/gou/rag"
+	"github.com/yaoapp/gou/rag/driver"
+	"github.com/yaoapp/yao/config"
+	neovision "github.com/yaoapp/yao/neo/vision"
+	vdriver "github.com/yaoapp/yao/neo/vision/driver"
+	"github.com/yaoapp/yao/test"
+)
+
+var (
+	// 1x1 transparent PNG for testing
+	testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+)
+
+func TestUpload(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ast := setupTestAssistant()
+	ctx := context.Background()
+
+	t.Run("Basic File Upload", func(t *testing.T) {
+		content := []byte("test content")
+		file := &multipart.FileHeader{
+			Filename: "test.txt",
+			Size:     int64(len(content)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "text/plain")
+
+		reader := bytes.NewReader(content)
+		fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
+			"sid":     "test-user",
+			"chat_id": "test-chat",
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, fileResp)
+		assert.Contains(t, fileResp.ID, "test-assistant/test-user/test-chat")
+		assert.Equal(t, len(content), fileResp.Bytes)
+		assert.Equal(t, "text/plain", fileResp.ContentType)
+	})
+
+	t.Run("File Size Limit", func(t *testing.T) {
+		content := make([]byte, MaxSize+1)
+		file := &multipart.FileHeader{
+			Filename: "large.txt",
+			Size:     int64(len(content)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "text/plain")
+
+		reader := bytes.NewReader(content)
+		_, err := ast.Upload(ctx, file, reader, nil)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "exceeds the maximum size")
+	})
+
+	t.Run("Invalid Content Type", func(t *testing.T) {
+		content := []byte("test")
+		file := &multipart.FileHeader{
+			Filename: "test.invalid",
+			Size:     int64(len(content)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "invalid/type")
+
+		reader := bytes.NewReader(content)
+		_, err := ast.Upload(ctx, file, reader, nil)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "not allowed")
+	})
+}
+
+func TestUploadWithRAG(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ast := setupTestAssistant()
+	ragEngine, ragUploader, ragVectorizer := setupTestRAG(t)
+	SetRAG(ragEngine, ragUploader, ragVectorizer, RAGSetting{IndexPrefix: "test_"})
+	defer func() {
+		rag = nil
+	}()
+	ctx := context.Background()
+
+	t.Run("Text File with RAG Enabled", func(t *testing.T) {
+		content := []byte("This is a test document for RAG indexing")
+		file := &multipart.FileHeader{
+			Filename: "test.txt",
+			Size:     int64(len(content)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "text/plain")
+
+		reader := bytes.NewReader(content)
+		fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
+			"sid":     "test-user",
+			"chat_id": "test-chat",
+			"rag":     true,
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, fileResp)
+		assert.NotEmpty(t, fileResp.DocIDs, "Document IDs should not be empty")
+
+		// Wait for indexing to complete
+		time.Sleep(500 * time.Millisecond)
+
+		// Verify the file was indexed
+		exists, err := ragEngine.HasDocument(ctx, "test_test-assistant-test-user-test-chat", fileResp.DocIDs[0])
+		assert.NoError(t, err)
+		assert.True(t, exists, "Document should exist in RAG index")
+	})
+
+	t.Run("Text File with RAG Disabled", func(t *testing.T) {
+		content := []byte("This is a test document with RAG disabled")
+		file := &multipart.FileHeader{
+			Filename: "test.txt",
+			Size:     int64(len(content)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "text/plain")
+
+		reader := bytes.NewReader(content)
+		fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
+			"sid":     "test-user",
+			"chat_id": "test-chat",
+			"rag":     false,
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, fileResp)
+		assert.Empty(t, fileResp.DocIDs, "Document IDs should be empty when RAG is disabled")
+	})
+}
+
+func setupTestAssistant() *Assistant {
+	ast := &Assistant{
+		ID:        "test-assistant",
+		Name:      "Test Assistant",
+		Connector: "test-connector",
+	}
+	return ast
+}
+
+func setupTestRAG(t *testing.T) (driver.Engine, driver.FileUpload, driver.Vectorizer) {
+	// Get test config
+	openaiKey := os.Getenv("OPENAI_API_KEY")
+	if openaiKey == "" {
+		t.Skip("OPENAI_API_KEY not set")
+	}
+
+	vectorizeConfig := driver.VectorizeConfig{
+		Model: os.Getenv("VECTORIZER_MODEL"),
+		Options: map[string]string{
+			"api_key": openaiKey,
+		},
+	}
+
+	// Qdrant config
+	host := os.Getenv("QDRANT_HOST")
+	if host == "" {
+		host = "localhost"
+	}
+
+	port := os.Getenv("QDRANT_PORT")
+	if port == "" {
+		port = "6334"
+	}
+
+	// Create vectorizer
+	vectorizer, err := gourag.NewVectorizer(gourag.DriverOpenAI, vectorizeConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create engine
+	engine, err := gourag.NewEngine(gourag.DriverQdrant, driver.IndexConfig{
+		Options: map[string]string{
+			"host":    host,
+			"port":    port,
+			"api_key": "",
+		},
+	}, vectorizer)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create file upload
+	fileUpload, err := gourag.NewFileUpload(gourag.DriverQdrant, engine, vectorizer)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return engine, fileUpload, vectorizer
+}
+
+func setupTestVision(t *testing.T) *neovision.Vision {
+	// Create test data directory
+	data, err := fs.Get("data")
+	assert.NoError(t, err)
+
+	// Write test image data
+	imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
+	assert.NoError(t, err)
+	_, err = data.WriteFile("/test.png", imgData, 0644)
+	assert.NoError(t, err)
+
+	cfg := &vdriver.Config{
+		Storage: vdriver.StorageConfig{
+			Driver: "local",
+			Options: map[string]interface{}{
+				"path":        "/__vision_test",
+				"compression": true,
+			},
+		},
+		Model: vdriver.ModelConfig{
+			Driver: "openai",
+			Options: map[string]interface{}{
+				"api_key": os.Getenv("OPENAI_API_KEY"),
+				"model":   os.Getenv("VISION_MODEL"),
+			},
+		},
+	}
+
+	v, err := neovision.New(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return v
+}
+
+func TestUploadWithVision(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ast := setupTestAssistant()
+	vision := setupTestVision(t)
+	SetVision(vision)
+	defer func() {
+		vision = nil
+	}()
+	ctx := context.Background()
+
+	t.Run("Image File with Vision Enabled", func(t *testing.T) {
+		imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
+		file := &multipart.FileHeader{
+			Filename: "test.png",
+			Size:     int64(len(imgData)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "image/png")
+
+		reader := bytes.NewReader(imgData)
+		fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
+			"vision": true,
+			"model":  "gpt-4-vision-preview",
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, fileResp)
+		if fileResp.URL == "" && fileResp.Description == "" {
+			t.Error("Either URL or Description should be set when vision is enabled")
+		}
+	})
+
+	t.Run("Image File with Vision Disabled", func(t *testing.T) {
+		imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
+		file := &multipart.FileHeader{
+			Filename: "test.png",
+			Size:     int64(len(imgData)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "image/png")
+
+		reader := bytes.NewReader(imgData)
+		fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
+			"vision": false,
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, fileResp)
+		assert.Empty(t, fileResp.URL, "Vision URL should be empty when vision is disabled")
+		assert.Empty(t, fileResp.Description, "Vision Description should be empty when vision is disabled")
+	})
+
+	t.Run("Image File with Non-Vision Model", func(t *testing.T) {
+		imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
+		file := &multipart.FileHeader{
+			Filename: "test.png",
+			Size:     int64(len(imgData)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "image/png")
+
+		reader := bytes.NewReader(imgData)
+		fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
+			"vision": true,
+			"model":  "gpt-4",
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, fileResp)
+		assert.Empty(t, fileResp.URL, "Vision URL should be empty for non-vision models")
+		assert.NotEmpty(t, fileResp.Description, "Vision Description should be set for non-vision models")
+	})
+}
+
+func TestDownload(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+
+	ast := setupTestAssistant()
+	ctx := context.Background()
+
+	t.Run("Download Existing File", func(t *testing.T) {
+		// First upload a file
+		content := []byte("test content")
+		file := &multipart.FileHeader{
+			Filename: "test.txt",
+			Size:     int64(len(content)),
+		}
+		file.Header = make(map[string][]string)
+		file.Header.Set("Content-Type", "text/plain")
+
+		reader := bytes.NewReader(content)
+		fileResp, err := ast.Upload(ctx, file, reader, nil)
+		assert.NoError(t, err)
+
+		// Then download it
+		downloadResp, err := ast.Download(ctx, fileResp.ID)
+		assert.NoError(t, err)
+		assert.NotNil(t, downloadResp)
+		assert.True(t, strings.HasPrefix(downloadResp.ContentType, "text/plain"), "Content-Type should start with text/plain")
+		assert.Equal(t, ".txt", downloadResp.Extension)
+
+		// Verify content
+		downloaded, err := io.ReadAll(downloadResp.Reader)
+		assert.NoError(t, err)
+		assert.Equal(t, content, downloaded)
+	})
+
+	t.Run("Download Non-Existent File", func(t *testing.T) {
+		_, err := ast.Download(ctx, "non-existent-file")
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "not found")
+	})
+}

neo/assistant/hooks.go

@@ -0,0 +1,229 @@
+package assistant
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	chatctx "github.com/yaoapp/yao/neo/context"
+	"github.com/yaoapp/yao/neo/message"
+)
+
+// HookInit initialize the assistant
+func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error) {
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, err := ast.call(ctx, "Init", context, input, c.Writer)
+	if err != nil {
+		if err.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	response := &ResHookInit{}
+	switch v := v.(type) {
+	case map[string]interface{}:
+		if res, ok := v["assistant_id"].(string); ok {
+			response.AssistantID = res
+		}
+		if res, ok := v["chat_id"].(string); ok {
+			response.ChatID = res
+		}
+
+		if res, ok := v["next"].(map[string]interface{}); ok {
+			response.Next = &NextAction{}
+			if name, ok := res["action"].(string); ok {
+				response.Next.Action = name
+			}
+			if payload, ok := res["payload"].(map[string]interface{}); ok {
+				response.Next.Payload = payload
+			}
+		}
+
+	case string:
+		response.AssistantID = v
+		response.ChatID = context.ChatID
+
+	case nil:
+		response.AssistantID = ast.ID
+		response.ChatID = context.ChatID
+	}
+
+	return response, nil
+}
+
+// HookStream Handle streaming response from LLM
+func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) {
+
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, err := ast.call(ctx, "Stream", context, input, output, c.Writer)
+	if err != nil {
+		if err.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	response := &ResHookStream{}
+	switch v := v.(type) {
+	case map[string]interface{}:
+		if res, ok := v["output"].(string); ok {
+			response.Output = res
+		}
+		if res, ok := v["next"].(map[string]interface{}); ok {
+			response.Next = &NextAction{}
+			if name, ok := res["action"].(string); ok {
+				response.Next.Action = name
+			}
+			if payload, ok := res["payload"].(map[string]interface{}); ok {
+				response.Next.Payload = payload
+			}
+		}
+
+		// Custom silent from hook
+		if res, ok := v["silent"].(bool); ok {
+			response.Silent = res
+		}
+
+	case string:
+		response.Output = v
+	}
+
+	return response, nil
+}
+
+// HookDone Handle completion of assistant response
+func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) {
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, err := ast.call(ctx, "Done", context, input, output, c.Writer)
+	if err != nil {
+		if err.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	response := &ResHookDone{
+		Input:  input,
+		Output: output,
+	}
+
+	switch v := v.(type) {
+	case map[string]interface{}:
+		if res, ok := v["output"].(string); ok {
+			response.Output = res
+		}
+		if res, ok := v["next"].(map[string]interface{}); ok {
+			response.Next = &NextAction{}
+			if name, ok := res["action"].(string); ok {
+				response.Next.Action = name
+			}
+			if payload, ok := res["payload"].(map[string]interface{}); ok {
+				response.Next.Payload = payload
+			}
+		}
+	case string:
+		response.Output = v
+	}
+
+	return response, nil
+}
+
+// HookFail Handle failure of assistant response
+func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, output string, err error) (*ResHookFail, error) {
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, callErr := ast.call(ctx, "Fail", context, input, output, err.Error(), c.Writer)
+	if callErr != nil {
+		if callErr.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, callErr
+	}
+
+	response := &ResHookFail{
+		Input:  input,
+		Output: output,
+		Error:  err.Error(),
+	}
+
+	switch v := v.(type) {
+	case map[string]interface{}:
+		if res, ok := v["output"].(string); ok {
+			response.Output = res
+		}
+		if res, ok := v["error"].(string); ok {
+			response.Error = res
+		}
+		if res, ok := v["next"].(map[string]interface{}); ok {
+			response.Next = &NextAction{}
+			if name, ok := res["action"].(string); ok {
+				response.Next.Action = name
+			}
+			if payload, ok := res["payload"].(map[string]interface{}); ok {
+				response.Next.Payload = payload
+			}
+		}
+	case string:
+		response.Output = v
+	}
+
+	return response, nil
+}
+
+// createTimeoutContext creates a timeout context with 5 seconds timeout
+func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second)
+	return ctx, cancel
+}
+
+// Call the script method
+func (ast *Assistant) call(ctx context.Context, method string, context chatctx.Context, args ...any) (interface{}, error) {
+	if ast.Script == nil {
+		return nil, nil
+	}
+
+	scriptCtx, err := ast.Script.NewContext(context.Sid, nil)
+	if err != nil {
+		return nil, err
+	}
+	defer scriptCtx.Close()
+
+	// Check if the method exists
+	if !scriptCtx.Global().Has(method) {
+		return nil, fmt.Errorf(HookErrorMethodNotFound)
+	}
+
+	// Create done channel for handling cancellation
+	done := make(chan struct{})
+	var result interface{}
+	var callErr error
+
+	go func() {
+		defer close(done)
+		// Call the method
+		args = append([]interface{}{context.Map()}, args...)
+		result, callErr = scriptCtx.Call(method, args...)
+	}()
+
+	// Wait for either context cancellation or method completion
+	select {
+	case <-ctx.Done():
+		scriptCtx.Close() // Force close the script context
+		return nil, ctx.Err()
+	case <-done:
+		return result, callErr
+	}
+}

neo/assistant/load_test.go

@@ -53,7 +53,7 @@ func TestLoad_LoadStore(t *testing.T) {
 				"assistant_id": "test-id",
 				"name":         "Test Assistant",
 				"avatar":       "test-avatar",
-				"connector":    "test-connector",
+				"connector":    "gpt-3_5-turbo",
 			},
 		},
 	}
@@ -67,7 +67,7 @@ func TestLoad_LoadStore(t *testing.T) {
 	assert.Equal(t, "test-id", assistant.ID)
 	assert.Equal(t, "Test Assistant", assistant.Name)
 	assert.Equal(t, "test-avatar", assistant.Avatar)
-	assert.Equal(t, "test-connector", assistant.Connector)
+	assert.Equal(t, "gpt-3_5-turbo", assistant.Connector)
 
 	// Test cache functionality
 	assistant2, err := LoadStore("test-id")

neo/assistant/types.go

@@ -5,17 +5,64 @@ import (
 	"io"
 	"mime/multipart"
 
+	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/gou/rag/driver"
 	v8 "github.com/yaoapp/gou/runtime/v8"
+	chatctx "github.com/yaoapp/yao/neo/context"
+	"github.com/yaoapp/yao/neo/message"
 	api "github.com/yaoapp/yao/openai"
 )
 
+const (
+	// HookErrorMethodNotFound is the error message for method not found
+	HookErrorMethodNotFound = "method not found"
+)
+
 // API the assistant API interface
 type API interface {
-	Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error
+	Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error
 	Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error)
 	Download(ctx context.Context, fileID string) (*FileResponse, error)
 	ReadBase64(ctx context.Context, fileID string) (string, error)
+	Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error
+	HookInit(c *gin.Context, ctx chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error)
+}
+
+// ResHookInit the response of the init hook
+type ResHookInit struct {
+	AssistantID string                 `json:"assistant_id,omitempty"`
+	ChatID      string                 `json:"chat_id,omitempty"`
+	Next        *NextAction            `json:"next,omitempty"`
+	Input       []message.Message      `json:"input,omitempty"`
+	Options     map[string]interface{} `json:"options,omitempty"`
+}
+
+// ResHookStream the response of the stream hook
+type ResHookStream struct {
+	Silent bool        `json:"silent,omitempty"` // Whether to suppress the output
+	Next   *NextAction `json:"next,omitempty"`   // The next action
+	Output string      `json:"output,omitempty"` // The output
+}
+
+// ResHookDone the response of the done hook
+type ResHookDone struct {
+	Next   *NextAction       `json:"next,omitempty"`
+	Input  []message.Message `json:"input,omitempty"`
+	Output string            `json:"output,omitempty"`
+}
+
+// ResHookFail the response of the fail hook
+type ResHookFail struct {
+	Next   *NextAction       `json:"next,omitempty"`
+	Input  []message.Message `json:"input,omitempty"`
+	Output string            `json:"output,omitempty"`
+	Error  string            `json:"error,omitempty"`
+}
+
+// NextAction the next action
+type NextAction struct {
+	Action  string                 `json:"action"`
+	Payload map[string]interface{} `json:"payload,omitempty"`
 }
 
 // RAG the RAG interface
@@ -70,6 +117,31 @@ type Assistant struct {
 	openai      *api.OpenAI              // OpenAI API
 }
 
+// VisionCapableModels list of LLM models that support vision capabilities
+var VisionCapableModels = map[string]bool{
+	// OpenAI Models
+	"gpt-4-vision-preview": true,
+	"gpt-4v":               true, // Alias for gpt-4-vision-preview
+
+	// Anthropic Models
+	"claude-3-opus":   true, // Most capable Claude model
+	"claude-3-sonnet": true, // Balanced Claude model
+	"claude-3-haiku":  true, // Fast and efficient Claude model
+
+	// Google Models
+	"gemini-pro-vision": true,
+
+	// Open Source Models
+	"llava-13b": true,
+	"cogvlm":    true,
+	"qwen-vl":   true,
+	"yi-vl":     true,
+
+	// Custom Models
+	"gpt-4o":      true, // Custom OpenAI compatible model
+	"gpt-4o-mini": true, // Custom OpenAI compatible model - mini version
+}
+
 // File the file
 type File struct {
 	ID          string   `json:"file_id"`
@@ -77,6 +149,9 @@ type File struct {
 	CreatedAt   int      `json:"created_at"`
 	Filename    string   `json:"filename"`
 	ContentType string   `json:"content_type"`
+	Description string   `json:"description,omitempty"` // Vision analysis result or other description
+	URL         string   `json:"url,omitempty"`         // Vision URL for vision-capable models
+	DocIDs      []string `json:"doc_ids,omitempty"`     // RAG document IDs
 }
 
 // FileResponse represents a file download response

neo/context.go

@@ -1,49 +0,0 @@
-package neo
-
-import (
-	"context"
-	"time"
-
-	jsoniter "github.com/json-iterator/go"
-	"github.com/yaoapp/kun/log"
-)
-
-// NewContext create a new context
-func NewContext(sid, cid, payload string) Context {
-	ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid}
-	if payload == "" {
-		return ctx
-	}
-
-	err := jsoniter.Unmarshal([]byte(payload), &ctx)
-	if err != nil {
-		log.Error("%s", err.Error())
-	}
-	return ctx
-}
-
-// NewContextWithCancel create a new context with cancel
-func NewContextWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
-	ctx := NewContext(sid, cid, payload)
-	return ContextWithCancel(ctx)
-}
-
-// NewContextWithTimeout create a new context with timeout
-func NewContextWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
-	ctx := NewContext(sid, cid, payload)
-	return ContextWithTimeout(ctx, timeout)
-}
-
-// ContextWithCancel create a new context
-func ContextWithCancel(parent Context) (Context, context.CancelFunc) {
-	new, cancel := context.WithCancel(parent.Context)
-	parent.Context = new
-	return parent, cancel
-}
-
-// ContextWithTimeout create a new context
-func ContextWithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
-	new, cancel := context.WithTimeout(parent.Context, timeout)
-	parent.Context = new
-	return parent, cancel
-}

neo/context/context.go

@@ -0,0 +1,122 @@
+package context
+
+import (
+	"context"
+	"time"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/yaoapp/kun/log"
+)
+
+// Context the context
+type Context struct {
+	context.Context
+	Sid         string                 `json:"sid" yaml:"-"`           // Session ID
+	ChatID      string                 `json:"chat_id,omitempty"`      // Chat ID, use to select chat
+	AssistantID string                 `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
+	Stack       string                 `json:"stack,omitempty"`
+	Path        string                 `json:"pathname,omitempty"`
+	FormData    map[string]interface{} `json:"formdata,omitempty"`
+	Field       *Field                 `json:"field,omitempty"`
+	Namespace   string                 `json:"namespace,omitempty"`
+	Config      map[string]interface{} `json:"config,omitempty"`
+	Signal      interface{}            `json:"signal,omitempty"`
+	Upload      *FileUpload            `json:"upload,omitempty"`
+}
+
+// Field the context field
+type Field struct {
+	Name     string                 `json:"name,omitempty"`
+	Type     string                 `json:"type,omitempty"`
+	Bind     string                 `json:"bind,omitempty"`
+	Props    map[string]interface{} `json:"props,omitempty"`
+	Children []interface{}          `json:"children,omitempty"`
+}
+
+// FileUpload the file upload
+type FileUpload struct {
+	Name     string `json:"name,omitempty"`
+	Type     string `json:"type,omitempty"`
+	Size     int64  `json:"size,omitempty"`
+	TempFile string `json:"temp_file,omitempty"`
+}
+
+// New create a new context
+func New(sid, cid, payload string) Context {
+	ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid}
+	if payload == "" {
+		return ctx
+	}
+
+	err := jsoniter.Unmarshal([]byte(payload), &ctx)
+	if err != nil {
+		log.Error("%s", err.Error())
+	}
+	return ctx
+}
+
+// NewWithCancel create a new context with cancel
+func NewWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
+	ctx := New(sid, cid, payload)
+	return WithCancel(ctx)
+}
+
+// NewWithTimeout create a new context with timeout
+func NewWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
+	ctx := New(sid, cid, payload)
+	return WithTimeout(ctx, timeout)
+}
+
+// WithCancel create a new context
+func WithCancel(parent Context) (Context, context.CancelFunc) {
+	new, cancel := context.WithCancel(parent.Context)
+	parent.Context = new
+	return parent, cancel
+}
+
+// WithTimeout create a new context
+func WithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
+	new, cancel := context.WithTimeout(parent.Context, timeout)
+	parent.Context = new
+	return parent, cancel
+}
+
+// Map the context to a map
+func (ctx *Context) Map() map[string]interface{} {
+	data := map[string]interface{}{
+		"sid": ctx.Sid,
+	}
+
+	if ctx.ChatID != "" {
+		data["chat_id"] = ctx.ChatID
+	}
+	if ctx.AssistantID != "" {
+		data["assistant_id"] = ctx.AssistantID
+	}
+	if ctx.Stack != "" {
+		data["stack"] = ctx.Stack
+	}
+	if ctx.Path != "" {
+		data["pathname"] = ctx.Path
+	}
+	if len(ctx.FormData) > 0 {
+		data["formdata"] = ctx.FormData
+	}
+	if ctx.Field != nil {
+		data["field"] = ctx.Field
+	}
+	if ctx.Namespace != "" {
+		data["namespace"] = ctx.Namespace
+	}
+	if len(ctx.Config) > 0 {
+		data["config"] = ctx.Config
+	}
+	if ctx.Signal != nil {
+		data["signal"] = ctx.Signal
+	}
+	if ctx.Upload != nil {
+		data["upload"] = ctx.Upload
+	}
+
+	return data
+}

neo/hooks.go

@@ -7,10 +7,11 @@ import (
 	"github.com/gin-gonic/gin"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/process"
+	chatctx "github.com/yaoapp/yao/neo/context"
 )
 
 // HookCreate create the assistant
-func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
+func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
 
 	// Default assistant
 	assistantID := neo.Use
@@ -69,7 +70,7 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi
 }
 
 // HookPrepare executes the prepare hook before AI is called
-func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
+func (neo *DSL) HookPrepare(ctx chatctx.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
 	if neo.Prepare == "" {
 		return messages, nil
 	}
@@ -114,7 +115,7 @@ func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]m
 }
 
 // HookWrite executes the write hook when response is received from AI
-func (neo *DSL) HookWrite(ctx Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) {
+func (neo *DSL) HookWrite(ctx chatctx.Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) {
 	if neo.Write == "" {
 		return []map[string]interface{}{response}, nil
 	}

neo/message/message.go

@@ -22,6 +22,8 @@ type Message struct {
 	IsDone      bool                   `json:"done,omitempty"`
 	Actions     []Action               `json:"actions,omitempty"`     // Conversation Actions for frontend
 	Attachments []Attachment           `json:"attachments,omitempty"` // File attachments
+	Role        string                 `json:"role,omitempty"`        // user, assistant, system ...
+	Name        string                 `json:"name,omitempty"`        // name for the message
 	Data        map[string]interface{} `json:"-"`
 }
 
@@ -134,12 +136,74 @@ func (m *Message) Error(message interface{}) *Message {
 	return m
 }
 
+// SetContent set the content
+func (m *Message) SetContent(content string) *Message {
+	if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
+		var msg Message
+		if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
+			m.Text = err.Error() + "\n" + content
+			return m
+		}
+		*m = msg
+	} else {
+		m.Text = content
+		m.Type = "text"
+	}
+	return m
+}
+
+// Content get the content
+func (m *Message) Content() string {
+	content := map[string]interface{}{"text": m.Text}
+	if m.Attachments != nil {
+		content["attachments"] = m.Attachments
+	}
+
+	if m.Type != "" {
+		content["type"] = m.Type
+	}
+	contentRaw, _ := jsoniter.MarshalToString(content)
+	return contentRaw
+}
+
+// ToMap convert to map
+func (m *Message) ToMap() map[string]interface{} {
+	return map[string]interface{}{
+		"content": m.Content(),
+		"role":    m.Role,
+		"name":    m.Name,
+	}
+}
+
 // Map set from map
 func (m *Message) Map(msg map[string]interface{}) *Message {
 	if msg == nil {
 		return m
 	}
 
+	// Content  {"text": "xxxx",  "attachments": ... }
+	if content, ok := msg["content"].(string); ok {
+		if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
+			var msg Message
+			if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
+				m.Text = err.Error() + "\n" + content
+				return m
+			}
+			*m = msg
+		} else {
+			m.Text = content
+			m.Type = "text"
+		}
+	}
+
+	if role, ok := msg["role"].(string); ok {
+		m.Role = role
+	}
+
+	if name, ok := msg["name"].(string); ok {
+		m.Name = name
+	}
+
 	if text, ok := msg["text"].(string); ok {
 		m.Text = text
 	}

neo/neo.go

@@ -4,41 +4,25 @@ import (
 	"fmt"
 	"os"
 	"strings"
-	"sync"
 
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/kun/log"
 	"github.com/yaoapp/yao/neo/assistant"
+	chatctx "github.com/yaoapp/yao/neo/context"
 	"github.com/yaoapp/yao/neo/message"
 )
 
-// Lock the assistant list
-var lock sync.Mutex = sync.Mutex{}
-
 // Answer reply the message
-func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error {
+func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error {
-	messages, err := neo.chatMessages(ctx, question)
+	var err error
-	if err != nil {
+	var ast assistant.API = neo.Assistant
-		msg := message.New().Error(err).Done()
+	if ctx.AssistantID != "" {
-		msg.Write(c.Writer)
+		ast, err = neo.Select(ctx.AssistantID)
-		return err
-	}
-
-	// Get the assistant_id, chat_id
-	res, err := neo.HookCreate(ctx, messages, c)
-	if err != nil {
-		msg := message.New().Error(err).Done()
-		msg.Write(c.Writer)
-		return err
-	}
-
-	// Select Assistant
-	ast, err := neo.Select(res.AssistantID)
 		if err != nil {
 			return err
 		}
-	// Chat with AI
+	}
-	return neo.chat(ast, ctx, messages, c)
+	return ast.Execute(c, ctx, question, nil)
 }
 
 // Select select an assistant
@@ -50,7 +34,7 @@ func (neo *DSL) Select(id string) (assistant.API, error) {
 }
 
 // GeneratePrompts generate prompts for the AI assistant
-func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) {
+func (neo *DSL) GeneratePrompts(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) {
 	prompts := `
 	Optimize the prompts for the AI assistant
 	1. Optimize prompts based on the user's input
@@ -68,13 +52,14 @@ func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silen
 }
 
 // GenerateChatTitle generate the chat title
-func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) {
+func (neo *DSL) GenerateChatTitle(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) {
 	prompts := `
 	Help me generate a title for the chat 
 	1. The title should be a short and concise description of the chat.
 	2. The title should be a single sentence.
 	3. The title should be in same language as the chat.
 	4. The title should be no more than 50 characters.
+	5. ANSWER ONLY THE TITLE CONTENT, FOR EXAMPLE: Chat with AI is a valid title, but "Chat with AI" is not a valid title.
 	`
 	isSilent := false
 	if len(silent) > 0 {
@@ -84,7 +69,7 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, sil
 }
 
 // GenerateWithAI generate content with AI, type can be "title", "prompts", etc.
-func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
+func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
 	messages := []map[string]interface{}{
 		{"role": "system", "content": systemPrompt},
 		{
@@ -119,7 +104,11 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy
 
 	// Chat with AI in background
 	go func() {
-		err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int {
+		msgList := make([]message.Message, len(messages))
+		for i, msg := range messages {
+			msgList[i] = *message.New().Map(msg)
+		}
+		err := ast.Chat(c.Request.Context(), msgList, neo.Option, func(data []byte) int {
 			select {
 			case <-clientBreak:
 				return 0 // break
@@ -187,7 +176,7 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy
 }
 
 // Upload upload a file
-func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
+func (neo *DSL) Upload(ctx chatctx.Context, c *gin.Context) (*assistant.File, error) {
 	// Get the file
 	tmpfile, err := c.FormFile("file")
 	if err != nil {
@@ -212,29 +201,30 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
 	}
 
 	// Get file info
-	ctx.Upload = &FileUpload{
+	ctx.Upload = &chatctx.FileUpload{
-		Bytes:       int(tmpfile.Size),
 		Name:     tmpfile.Filename,
-		ContentType: tmpfile.Header.Get("Content-Type"),
+		Type:     tmpfile.Header.Get("Content-Type"),
-		Option:      option,
+		Size:     tmpfile.Size,
+		TempFile: tmpfile.Filename,
 	}
 
-	res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c)
+	// Default use the assistant in context
+	ast := neo.Assistant
+	if ctx.ChatID == "" {
+		if ctx.AssistantID == "" {
+			return nil, fmt.Errorf("assistant_id is required")
+		}
+		ast, err = neo.Select(ctx.AssistantID)
 		if err != nil {
 			return nil, err
 		}
-
-	// Select Assistant
-	ast, err := neo.Select(res.AssistantID)
-	if err != nil {
-		return nil, err
 	}
 
 	return ast.Upload(ctx, tmpfile, reader, option)
 }
 
 // Download downloads a file
-func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, error) {
+func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileResponse, error) {
 	// Get file_id from query string
 	fileID := c.Query("file_id")
 	if fileID == "" {
@@ -256,133 +246,3 @@ func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse,
 	// Download file using the assistant
 	return ast.Download(ctx.Context, fileID)
 }
-
-// chat chat with AI
-func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error {
-	if ast == nil {
-		msg := message.New().Error("assistant is not initialized").Done()
-		msg.Write(c.Writer)
-		return fmt.Errorf("assistant is not initialized")
-	}
-
-	clientBreak := make(chan bool, 1)
-	done := make(chan bool, 1)
-	content := []byte{}
-
-	// Chat with AI in background
-	go func() {
-		err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int {
-			select {
-			case <-clientBreak:
-				return 0 // break
-
-			default:
-				msg := message.NewOpenAI(data)
-				if msg == nil {
-					return 1 // continue
-				}
-
-				// Handle error
-				if msg.Type == "error" {
-					value := msg.String()
-					message.New().Error(value).Done().Write(c.Writer)
-					return 0 // break
-				}
-
-				// Append content and send message
-				content = msg.Append(content)
-				value := msg.String()
-				if value != "" {
-					message.New().
-						Map(map[string]interface{}{
-							"text": value,
-							"done": msg.IsDone,
-						}).
-						Write(c.Writer)
-				}
-
-				// Complete the stream
-				if msg.IsDone {
-					if value == "" {
-						msg.Write(c.Writer)
-					}
-					done <- true
-					return 0 // break
-				}
-
-				return 1 // continue
-			}
-		})
-
-		if err != nil {
-			log.Error("Chat error: %s", err.Error())
-			message.New().Error(err).Done().Write(c.Writer)
-		}
-
-		// Save chat history
-		if len(content) > 0 {
-			neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages)
-		}
-
-		done <- true
-	}()
-
-	// Wait for completion or client disconnect
-	select {
-	case <-done:
-		return nil
-	case <-c.Writer.CloseNotify():
-		clientBreak <- true
-		return nil
-	}
-}
-
-// chatMessages get the chat messages
-func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) {
-
-	history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID)
-	if err != nil {
-		return nil, err
-	}
-
-	messages := []map[string]interface{}{}
-	messages = append(messages, history...)
-	if len(content) == 0 {
-		return messages, nil
-	}
-
-	// Add user message
-	messages = append(messages, map[string]interface{}{"role": "user", "content": content[0], "name": ctx.Sid})
-	return messages, nil
-}
-
-// saveHistory save the history
-func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages []map[string]interface{}) {
-
-	if len(content) > 0 && sid != "" && len(messages) > 0 {
-		err := neo.Store.SaveHistory(
-			sid,
-			[]map[string]interface{}{
-				{"role": "user", "content": messages[len(messages)-1]["content"], "name": sid},
-				{"role": "assistant", "content": string(content), "name": sid},
-			},
-			chatID,
-			nil,
-		)
-
-		if err != nil {
-			log.Error("Save history error: %s", err.Error())
-		}
-	}
-}
-
-// sendMessage sends a message to the client
-func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error {
-	if msg, ok := data.(map[string]interface{}); ok {
-		if !message.New().Map(msg).Write(w) {
-			return fmt.Errorf("failed to write message to stream")
-		}
-		return nil
-	}
-	return fmt.Errorf("invalid message data type")
-}

neo/types.go

@@ -1,8 +1,6 @@
 package neo
 
 import (
-	"context"
-
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/yao/neo/assistant"
 	"github.com/yaoapp/yao/neo/rag"
@@ -48,22 +46,6 @@ type Mention struct {
 	Type   string `json:"type,omitempty"`
 }
 
-// Context the context
-type Context struct {
-	Sid             string                 `json:"sid" yaml:"-"`           // Session ID
-	ChatID          string                 `json:"chat_id,omitempty"`      // Chat ID, use to select chat
-	AssistantID     string                 `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
-	Stack           string                 `json:"stack,omitempty"`
-	Path            string                 `json:"pathname,omitempty"`
-	FormData        map[string]interface{} `json:"formdata,omitempty"`
-	Field           *Field                 `json:"field,omitempty"`
-	Namespace       string                 `json:"namespace,omitempty"`
-	Config          map[string]interface{} `json:"config,omitempty"`
-	Signal          interface{}            `json:"signal,omitempty"`
-	Upload          *FileUpload            `json:"upload,omitempty"`
-	context.Context `json:"-" yaml:"-"`
-}
-
 // Field the context field
 type Field struct {
 	Name string `json:"name,omitempty"`

neo/vision/driver/openai/model.go

@@ -52,11 +52,17 @@ func New(options map[string]interface{}) (*Model, error) {
 }
 
 // Analyze analyze image using OpenAI vision model
-func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error) {
+func (model *Model) Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error) {
 	if model.APIKey == "" {
 		return nil, fmt.Errorf("api_key is required")
 	}
 
+	// Use default prompt if none provided
+	userPrompt := model.Prompt
+	if len(prompt) > 0 && prompt[0] != "" {
+		userPrompt = prompt[0]
+	}
+
 	// Check if fileID is a URL or base64 data
 	var imageURL string
 	if strings.HasPrefix(fileID, "data:image/") {
@@ -103,7 +109,7 @@ func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) (
 				"content": []map[string]interface{}{
 					{
 						"type": "text",
-						"text": prompt,
+						"text": userPrompt,
 					},
 					{
 						"type": "image_url",

neo/vision/driver/openai/model_test.go

@@ -146,4 +146,49 @@ func TestOpenAIModel(t *testing.T) {
 		assert.Error(t, err)
 		assert.Contains(t, err.Error(), "OpenAI API error")
 	})
+
+	t.Run("Analyze with Default Prompt", func(t *testing.T) {
+		model, err := New(map[string]interface{}{
+			"api_key": os.Getenv("OPENAI_API_KEY"),
+			"model":   os.Getenv("VISION_MODEL"),
+			"prompt":  "Default test prompt",
+		})
+		assert.NoError(t, err)
+
+		// Use base64 image data without providing a prompt
+		result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64)
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+		assert.NotEmpty(t, result["description"])
+	})
+
+	t.Run("Analyze with Custom Prompt Overriding Default", func(t *testing.T) {
+		model, err := New(map[string]interface{}{
+			"api_key": os.Getenv("OPENAI_API_KEY"),
+			"model":   os.Getenv("VISION_MODEL"),
+			"prompt":  "Default test prompt",
+		})
+		assert.NoError(t, err)
+
+		// Use base64 image data with custom prompt
+		result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt")
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+		assert.NotEmpty(t, result["description"])
+	})
+
+	t.Run("Analyze with Empty Custom Prompt", func(t *testing.T) {
+		model, err := New(map[string]interface{}{
+			"api_key": os.Getenv("OPENAI_API_KEY"),
+			"model":   os.Getenv("VISION_MODEL"),
+			"prompt":  "Default test prompt",
+		})
+		assert.NoError(t, err)
+
+		// Use base64 image data with empty prompt (should use default)
+		result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "")
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+		assert.NotEmpty(t, result["description"])
+	})
 }

neo/vision/driver/types.go

@@ -32,7 +32,9 @@ type Storage interface {
 
 // Model the vision model interface
 type Model interface {
-	Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error)
+	// Analyze analyzes an image file
+	// If prompt is empty, it will use the default prompt from model.options.prompt
+	Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error)
 }
 
 // Response the vision response

neo/vision/vision.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"os"
 	"strings"
 	"time"
 
@@ -13,6 +14,30 @@ import (
 	"github.com/yaoapp/yao/neo/vision/driver/s3"
 )
 
+// parseEnvValue parse environment variable if the value starts with $ENV.
+func parseEnvValue(value string) string {
+	if strings.HasPrefix(value, "$ENV.") {
+		envKey := strings.TrimPrefix(value, "$ENV.")
+		if envVal := os.Getenv(envKey); envVal != "" {
+			return envVal
+		}
+	}
+	return value
+}
+
+// convertOptions convert interface{} options map to string map and parse environment variables
+func convertOptions(options map[string]interface{}) map[string]interface{} {
+	converted := make(map[string]interface{})
+	for k, v := range options {
+		if str, ok := v.(string); ok {
+			converted[k] = parseEnvValue(str)
+		} else {
+			converted[k] = v
+		}
+	}
+	return converted
+}
+
 // Vision the vision service
 type Vision struct {
 	storage driver.Storage
@@ -22,20 +47,24 @@ type Vision struct {
 // New create a new vision service
 func New(cfg *driver.Config) (*Vision, error) {
 
+	// Parse environment variables in options
+	storageOptions := convertOptions(cfg.Storage.Options)
+	modelOptions := convertOptions(cfg.Model.Options)
+
 	// Create storage driver
 	var storage driver.Storage
 	var err error
 	switch cfg.Storage.Driver {
 	case "local":
-		storage, err = local.New(cfg.Storage.Options)
+		storage, err = local.New(storageOptions)
 	case "s3":
 		// Convert expiration string to duration if present
-		if exp, ok := cfg.Storage.Options["expiration"].(string); ok {
+		if exp, ok := storageOptions["expiration"].(string); ok {
 			if duration, err := time.ParseDuration(exp); err == nil {
-				cfg.Storage.Options["expiration"] = duration
+				storageOptions["expiration"] = duration
 			}
 		}
-		storage, err = s3.New(cfg.Storage.Options)
+		storage, err = s3.New(storageOptions)
 	default:
 		return nil, fmt.Errorf("storage driver %s not supported", cfg.Storage.Driver)
 	}
@@ -47,7 +76,7 @@ func New(cfg *driver.Config) (*Vision, error) {
 	var model driver.Model
 	switch cfg.Model.Driver {
 	case "openai":
-		model, err = openai.New(cfg.Model.Options)
+		model, err = openai.New(modelOptions)
 	default:
 		return nil, fmt.Errorf("model driver %s not supported", cfg.Model.Driver)
 	}
@@ -75,7 +104,7 @@ func (v *Vision) Upload(ctx context.Context, filename string, reader io.Reader,
 }
 
 // Analyze analyze image using vision model
-func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*driver.Response, error) {
+func (v *Vision) Analyze(ctx context.Context, fileID string, prompt ...string) (*driver.Response, error) {
 	if v.model == nil {
 		return nil, fmt.Errorf("model is required")
 	}
@@ -92,7 +121,7 @@ func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*dr
 		}
 	}
 
-	result, err := v.model.Analyze(ctx, url, prompt)
+	result, err := v.model.Analyze(ctx, url, prompt...)
 	if err != nil {
 		return nil, err
 	}

neo/vision/vision_test.go

@@ -331,6 +331,96 @@ func TestVision(t *testing.T) {
 		assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
 		assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
 	})
+
+	t.Run("Analyze Image with Default Prompt", func(t *testing.T) {
+		// Create vision service with default prompt
+		cfg := &driver.Config{
+			Storage: driver.StorageConfig{
+				Driver: "local",
+				Options: map[string]interface{}{
+					"path":        "/__vision_test",
+					"compression": true,
+				},
+			},
+			Model: driver.ModelConfig{
+				Driver: "openai",
+				Options: map[string]interface{}{
+					"api_key": os.Getenv("OPENAI_API_KEY"),
+					"model":   os.Getenv("VISION_MODEL"),
+					"prompt":  "Default test prompt",
+				},
+			},
+		}
+
+		vision, err := New(cfg)
+		assert.NoError(t, err)
+
+		// Use base64 data without providing a prompt
+		result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64)
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+		assert.NotEmpty(t, result.Description)
+	})
+
+	t.Run("Analyze Image with Custom Prompt", func(t *testing.T) {
+		// Create vision service with default prompt
+		cfg := &driver.Config{
+			Storage: driver.StorageConfig{
+				Driver: "local",
+				Options: map[string]interface{}{
+					"path":        "/__vision_test",
+					"compression": true,
+				},
+			},
+			Model: driver.ModelConfig{
+				Driver: "openai",
+				Options: map[string]interface{}{
+					"api_key": os.Getenv("OPENAI_API_KEY"),
+					"model":   os.Getenv("VISION_MODEL"),
+					"prompt":  "Default test prompt",
+				},
+			},
+		}
+
+		vision, err := New(cfg)
+		assert.NoError(t, err)
+
+		// Use base64 data with custom prompt
+		result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt")
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+		assert.NotEmpty(t, result.Description)
+	})
+
+	t.Run("Analyze Image with Empty Custom Prompt", func(t *testing.T) {
+		// Create vision service with default prompt
+		cfg := &driver.Config{
+			Storage: driver.StorageConfig{
+				Driver: "local",
+				Options: map[string]interface{}{
+					"path":        "/__vision_test",
+					"compression": true,
+				},
+			},
+			Model: driver.ModelConfig{
+				Driver: "openai",
+				Options: map[string]interface{}{
+					"api_key": os.Getenv("OPENAI_API_KEY"),
+					"model":   os.Getenv("VISION_MODEL"),
+					"prompt":  "Default test prompt",
+				},
+			},
+		}
+
+		vision, err := New(cfg)
+		assert.NoError(t, err)
+
+		// Use base64 data with empty prompt (should use default)
+		result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "")
+		assert.NoError(t, err)
+		assert.NotNil(t, result)
+		assert.NotEmpty(t, result.Description)
+	})
 }
 
 func createTestVision(baseURL string) (*Vision, error) {