diff --git a/neo/api.go b/neo/api.go index 7084e764..8e906eb1 100644 --- a/neo/api.go +++ b/neo/api.go @@ -5,7 +5,6 @@ import ( "io" "net/url" "os" - "path/filepath" "strconv" "strings" "time" @@ -302,6 +301,105 @@ func (neo *DSL) handleUpload(c *gin.Context) { c.Done() } +// handleDownload handles the download request +func (neo *DSL) handleDownload(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + c.JSON(400, gin.H{"message": "sid is required", "code": 400}) + c.Done() + return + } + + uid, _, err := neo.UserOrGuestID(sid) + if err != nil { + c.JSON(401, gin.H{"message": fmt.Sprintf("Unauthorized, %s", err.Error()), "code": 401}) + c.Done() + return + } + + if uid == nil || uid == "" { + c.JSON(401, gin.H{"message": "Unauthorized", "code": 401}) + c.Done() + return + } + + fileID := c.Query("file_id") + if fileID == "" { + c.JSON(400, gin.H{"message": "file_id is required", "code": 400}) + c.Done() + return + } + + // Get the attachment + attach, err := neo.Store.GetAttachment(fileID) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + // Validate the permission ( Will be supported scope validation in the future ) + if (attach["public"] == 0 || attach["public"] == false) && attach["uid"] != uid { + c.JSON(403, gin.H{"message": "Forbidden", "code": 403}) + c.Done() + return + } + + storage, ok := attach["manager"].(string) + if !ok { + c.JSON(400, gin.H{"message": "Invalid storage", "code": 400}) + c.Done() + return + } + + // Get the manager + manager, ok := attachment.Managers[storage] + if !ok { + c.JSON(400, gin.H{"message": "Invalid storage", "code": 400}) + c.Done() + return + } + + name, ok := attach["name"].(string) + if !ok { + c.JSON(400, gin.H{"message": "Invalid name", "code": 400}) + c.Done() + return + } + + name = strings.TrimSuffix(name, ".gz") + contentType, ok := attach["content_type"].(string) + if !ok { + c.JSON(400, gin.H{"message": "Invalid content type", "code": 400}) + c.Done() + return + } + + handle, err := manager.Download(c.Request.Context(), fileID) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + defer handle.Reader.Close() + + // Set the response headers + encoded := url.PathEscape(name) + disposition := fmt.Sprintf(`attachment; filename="%s"`, encoded) + c.Header("Content-Type", contentType) + c.Header("Content-Disposition", disposition) + + // Copy the file content to response + _, err = io.Copy(c.Writer, handle.Reader) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + return + } + c.Done() + return + +} + // handleChat handles the chat request func (neo *DSL) handleChat(c *gin.Context) { // Set headers for SSE @@ -431,49 +529,6 @@ func (neo *DSL) handleChatHistory(c *gin.Context) { c.Done() } -// handleDownload handles the download request -func (neo *DSL) handleDownload(c *gin.Context) { - sid := c.GetString("__sid") - if sid == "" { - c.JSON(400, gin.H{"message": "sid is required", "code": 400}) - c.Done() - return - } - - fileID := c.Query("file_id") - if fileID == "" { - c.JSON(400, gin.H{"message": "file_id is required", "code": 400}) - c.Done() - return - } - - // Set the context - ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "") - defer cancel() - - // Download the file - fileResponse, err := neo.Download(ctx, c) - if err != nil { - c.JSON(500, gin.H{"message": err.Error(), "code": 500}) - c.Done() - return - } - defer fileResponse.Reader.Close() - - // Set response headers - c.Header("Content-Type", fileResponse.ContentType) - if disposition := c.Query("disposition"); disposition == "attachment" { - c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filepath.Base(fileID)+fileResponse.Extension)) - } - - // Copy the file content to response - _, err = io.Copy(c.Writer, fileResponse.Reader) - if err != nil { - c.JSON(500, gin.H{"message": err.Error(), "code": 500}) - return - } -} - // getCorsHandlers returns CORS middleware handlers func (neo *DSL) getCorsHandlers() ([]gin.HandlerFunc, error) { if len(neo.Allows) == 0 { @@ -511,8 +566,9 @@ func (neo *DSL) corsMiddleware(allowsMap map[string]bool) gin.HandlerFunc { // Set CORS headers c.Header("Access-Control-Allow-Origin", origin) c.Header("Access-Control-Allow-Credentials", "true") - c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range") + c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Disposition, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range") c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") + c.Header("Access-Control-Expose-Headers", "Content-Type, Content-Disposition, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range") if c.Request.Method == "OPTIONS" { c.AbortWithStatus(204) @@ -529,9 +585,10 @@ func (neo *DSL) optionsHandler(c *gin.Context) { if origin != "" { c.Header("Access-Control-Allow-Origin", origin) c.Header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS") - c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization, Accept, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range") + c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Disposition, Authorization, Accept, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range") c.Header("Access-Control-Allow-Credentials", "true") c.Header("Access-Control-Max-Age", "86400") // 24 hours + c.Header("Access-Control-Expose-Headers", "Content-Type, Content-Disposition, Authorization, Accept, Content-Sync, Content-Fingerprint, Content-Uid, Content-Range") } c.AbortWithStatus(204) } diff --git a/neo/assistant/attachment.go b/neo/assistant/attachment.go deleted file mode 100644 index f018d21f..00000000 --- a/neo/assistant/attachment.go +++ /dev/null @@ -1,331 +0,0 @@ -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 - } - - handleVision := false - if vv, has := option["vision"]; has { - switch v := vv.(type) { - case bool: - handleVision = v - case string: - handleVision = v == "true" || v == "1" || v == "yes" || v == "on" || v == "enable" - } - } - - if !handleVision { - return nil - } - - // Check if file is an image - if !strings.HasPrefix(file.ContentType, "image/") { - return nil - } - - // 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()) - } - - // The model is vision capable - if ast.vision { - // 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 - return nil - } - - // 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["description"].(string); ok { - file.Description = desc - } else 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 -} diff --git a/neo/assistant/attachment_test.go b/neo/assistant/attachment_test.go deleted file mode 100644 index 17a7d3b8..00000000 --- a/neo/assistant/attachment_test.go +++ /dev/null @@ -1,363 +0,0 @@ -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") - }) -} diff --git a/neo/neo.go b/neo/neo.go index 9e6ff66e..2f63593c 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -1,8 +1,6 @@ package neo import ( - "fmt" - "github.com/gin-gonic/gin" "github.com/yaoapp/gou/session" "github.com/yaoapp/yao/neo/assistant" @@ -61,29 +59,3 @@ func (neo *DSL) UserOrGuestID(sid string) (interface{}, bool, error) { } return userID, false, nil } - -// Download downloads a file -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 == "" { - return nil, fmt.Errorf("file_id is required") - } - - // Get assistant_id from context or query - // res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c) - // if err != nil { - // return nil, err - // } - - // Select Assistant - ast, err := neo.Select(neo.Use.Default) - if err != nil { - return nil, err - } - - // Download file using the assistant - // return ast.Download(ctx.Context, fileID) - fmt.Println(ast) - return nil, nil -}