Enhance file upload handling and deduplication features in neo package

- Updated UserOrGuestID method to return guest status alongside user ID. - Introduced UploadOption struct to encapsulate additional upload parameters, including public visibility and collection ID. - Enhanced handleUpload method to validate collection ID for knowledge storage and improved error handling. - Updated GetHeader function to include file size in the header processing. - Implemented content fingerprinting for file deduplication and added support for upload status tracking. - Enhanced README documentation to reflect new features and usage examples for file uploads and deduplication.
2025-06-01 17:46:28 +08:00 · 2025-06-01 17:46:28 +08:00 · e6a0f37a66
commit e6a0f37a66
parent 68e4e64e35
7 changed files with 294 additions and 27 deletions
--- a/neo/api.go
+++ b/neo/api.go
@ -179,7 +179,7 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 		sid = uuid.New().String()
 	}

-	uid, err := neo.UserOrGuestID(sid)
+	uid, isGuest, err := neo.UserOrGuestID(sid)
 	if err != nil {
 		c.JSON(401, gin.H{"message": fmt.Sprintf("Unauthorized, %s", err.Error()), "code": 401})
 		c.Done()
@ -209,7 +209,7 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 	}

 	// Get Option from form data
-	var option attachment.UploadOption
+	var option UploadOption
 	err = c.ShouldBind(&option)
 	if err != nil {
 		c.JSON(400, gin.H{"message": err.Error(), "code": 400})
@ -226,6 +226,13 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 			c.Done()
 			return
 		}
+
+	} else if storage == "knowledge" {
+		if option.CollectionID == "" {
+			c.JSON(400, gin.H{"message": "collection_id is required", "code": 400})
+			c.Done()
+			return
+		}
 	}

 	// Get the file
@ -239,7 +246,7 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 	// Open the file
 	reader, err := file.Open()
 	if err != nil {
-		c.JSON(400, gin.H{"message": err.Error(), "code": 400})
+		c.JSON(500, gin.H{"message": err.Error(), "code": 500})
 		c.Done()
 		return
 	}
@ -249,13 +256,48 @@ func (neo *DSL) handleUpload(c *gin.Context) {
 	}()

 	// Upload the file
-	header := attachment.GetHeader(c.Request.Header, file.Header)
-	res, err := manager.Upload(c.Request.Context(), header, reader, option)
+	header := attachment.GetHeader(c.Request.Header, file.Header, file.Size)
+	res, err := manager.Upload(c.Request.Context(), header, reader, option.UploadOption)
 	if err != nil {
-		c.JSON(400, gin.H{"message": err.Error(), "code": 500})
+		c.JSON(500, gin.H{"message": err.Error(), "code": 500})
 		c.Done()
 		return
 	}
+
+	// if storage is chat or knowledge, save the file to the store
+	if storage == "chat" || storage == "knowledge" {
+
+		attachment := map[string]interface{}{
+			"file_id":      res.ID,
+			"uid":          uid,
+			"guest":        isGuest,
+			"manager":      storage,
+			"public":       option.Public,
+			"name":         option.OriginalFilename,
+			"content_type": res.ContentType,
+			"bytes":        res.Bytes,
+			"gzip":         option.Gzip,
+			"status":       res.Status,
+		}
+
+		// Set the scope
+		if option.Scope != nil {
+			attachment["scope"] = option.Scope
+		}
+
+		// Set the collection_id
+		if option.CollectionID != "" {
+			attachment["collection_id"] = option.CollectionID
+		}
+
+		_, err = neo.Store.SaveAttachment(attachment)
+		if err != nil {
+			c.JSON(500, gin.H{"message": err.Error(), "code": 500})
+			c.Done()
+			return
+		}
+	}
+
 	c.JSON(200, map[string]interface{}{"data": res})
 	c.Done()
 }
@ -469,7 +511,7 @@ 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-Uid, Content-Range")
+		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-Methods", "GET, POST, OPTIONS")

 		if c.Request.Method == "OPTIONS" {
@ -487,7 +529,7 @@ 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-Uid, Content-Range")
+		c.Header("Access-Control-Allow-Headers", "Content-Type, 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
 	}
--- a/neo/attachment/README.md
+++ b/neo/attachment/README.md
@ -6,6 +6,7 @@ A comprehensive file upload package for Go that supports chunked uploads, file f

 - **Multiple Storage Backends**: Local filesystem and S3-compatible storage
 - **Chunked Upload Support**: Handle large files with standard HTTP Content-Range headers
+- **File Deduplication**: Content-based fingerprinting to avoid duplicate uploads
 - **File Compression**:
  - Gzip compression for any file type
  - Image compression with configurable size limits
@ -16,6 +17,8 @@ A comprehensive file upload package for Go that supports chunked uploads, file f
 - **Flexible File Organization**: Hierarchical storage with user/chat/assistant organization
 - **Multiple Read Methods**: Stream, bytes, and base64 encoding
 - **Global Manager Registry**: Support for registering and accessing managers globally
+- **Upload Status Tracking**: Track upload progress with status field
+- **Content Synchronization**: Support for synchronized uploads with Content-Sync header

 ## Installation

@ -38,8 +41,14 @@ import (
 )

 func main() {
-    // Create a manager
-    manager, err := attachment.New(attachment.ManagerOption{
+    // Create a manager with default settings
+    manager, err := attachment.RegisterDefault("uploads")
+    if err != nil {
+        panic(err)
+    }
+
+    // Or create a custom manager
+    customManager, err := attachment.New(attachment.ManagerOption{
        Driver:       "local",
        MaxSize:      "20M",
        ChunkSize:    "2M",
@ -64,8 +73,9 @@ func main() {
    fileHeader.Header.Set("Content-Type", "text/plain")

    option := attachment.UploadOption{
-        UserID: "user123",
-        ChatID: "chat456",
+        UserID:           "user123",
+        ChatID:           "chat456",
+        OriginalFilename: "my_document.txt", // Preserve original filename
    }

    file, err := manager.Upload(context.Background(), fileHeader, strings.NewReader(content), option)
@ -73,6 +83,11 @@ func main() {
        panic(err)
    }

+    // Check upload status
+    if file.Status == "uploaded" {
+        fmt.Printf("File uploaded successfully: %s\n", file.ID)
+    }
+
    // Read the file back
    data, err := manager.Read(context.Background(), file.ID)
    if err != nil {
@ -247,7 +262,10 @@ if err != nil {
 You can register managers globally for easy access:

 ```go
-// Register managers
+// Register default manager with sensible defaults
+attachment.RegisterDefault("main")
+
+// Register custom managers
 attachment.Register("local", "local", attachment.ManagerOption{
    Driver: "local",
    Options: map[string]interface{}{
@ -267,6 +285,7 @@ attachment.Register("s3", "s3", attachment.ManagerOption{
 // Use global managers
 localManager := attachment.Managers["local"]
 s3Manager := attachment.Managers["s3"]
+defaultManager := attachment.Managers["main"]
 ```

 ## File Organization
@ -355,7 +374,9 @@ Options for file upload:
 - `CompressImage`: Enable image compression
 - `CompressSize`: Maximum image dimension (default: 1920)
 - `Gzip`: Enable gzip compression
+- `Knowledge`: Push to knowledge base
 - `UserID`, `ChatID`, `AssistantID`: Organization IDs
+- `OriginalFilename`: Original filename to preserve (avoids encoding issues)

 #### `File`

@ -366,6 +387,7 @@ Uploaded file information:
 - `ContentType`: MIME type
 - `Bytes`: File size
 - `CreatedAt`: Upload timestamp
+- `Status`: Upload status ("uploading", "uploaded", "indexing", "indexed", "upload_failed", "index_failed")

 #### `FileResponse`

@ -448,3 +470,175 @@ The package includes comprehensive tests for:
 ## License

 This package is part of the Yao project and follows the same license terms.
+
+### File Deduplication with Fingerprints
+
+The package supports file deduplication using content fingerprints:
+
+```go
+// Set a content fingerprint to enable deduplication
+fileHeader := &attachment.FileHeader{
+    FileHeader: &multipart.FileHeader{
+        Filename: "document.pdf",
+        Size:     fileSize,
+        Header:   make(map[string][]string),
+    },
+}
+fileHeader.Header.Set("Content-Type", "application/pdf")
+fileHeader.Header.Set("Content-Fingerprint", "sha256:abcdef123456") // Content-based hash
+
+file, err := manager.Upload(ctx, fileHeader, reader, option)
+```
+
+### Content Synchronization
+
+For synchronized uploads across multiple clients:
+
+```go
+// Enable content synchronization
+fileHeader.Header.Set("Content-Sync", "true")
+
+// Each client can upload the same content with the same fingerprint
+// The system will deduplicate based on the content fingerprint
+```
+
+### Chunked Upload with Enhanced Headers
+
+For large files, you can upload in chunks using standard HTTP Content-Range headers with additional metadata:
+
+```go
+// Upload chunks with unique identifier and fingerprint
+totalSize := int64(1024000) // 1MB file
+chunkSize := int64(1024)    // 1KB chunks
+uid := "unique-file-id-123"
+fingerprint := "sha256:content-hash-here"
+
+for start := int64(0); start < totalSize; start += chunkSize {
+    end := start + chunkSize - 1
+    if end >= totalSize {
+        end = totalSize - 1
+    }
+
+    chunkData := make([]byte, end-start+1)
+    // ... fill chunkData with actual data ...
+
+    chunkHeader := &attachment.FileHeader{
+        FileHeader: &multipart.FileHeader{
+            Filename: "large_file.zip",
+            Size:     end - start + 1,
+            Header:   make(map[string][]string),
+        },
+    }
+    chunkHeader.Header.Set("Content-Type", "application/zip")
+    chunkHeader.Header.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, totalSize))
+    chunkHeader.Header.Set("Content-Uid", uid)
+    chunkHeader.Header.Set("Content-Fingerprint", fingerprint)
+    chunkHeader.Header.Set("Content-Sync", "true") // Enable synchronization
+
+    option := attachment.UploadOption{
+        UserID:           "user123",
+        ChatID:           "chat456",
+        OriginalFilename: "my_large_file.zip", // Preserve original name
+    }
+
+    file, err := manager.Upload(ctx, chunkHeader, bytes.NewReader(chunkData), option)
+    if err != nil {
+        return err
+    }
+
+    // Check if upload is complete
+    if file.Status == "uploaded" {
+        fmt.Printf("Upload complete: %s\n", file.ID)
+        break
+    } else if file.Status == "uploading" {
+        fmt.Printf("Chunk uploaded, progress: %d/%d\n", chunkHeader.GetChunkSize(), chunkHeader.GetTotalSize())
+    }
+}
+```
+
+### FileHeader Methods
+
+The `FileHeader` type provides several utility methods:
+
+```go
+// Get unique identifier for chunked uploads
+uid := fileHeader.UID()
+
+// Get content fingerprint for deduplication
+fingerprint := fileHeader.Fingerprint()
+
+// Get byte range for chunked uploads
+rangeHeader := fileHeader.Range()
+
+// Check if synchronization is enabled
+isSync := fileHeader.Sync()
+
+// Check if this is a chunked upload
+isChunk := fileHeader.IsChunk()
+
+// Check if upload is complete (for chunked uploads)
+isComplete := fileHeader.Complete()
+
+// Get detailed chunk information
+start, end, total, err := fileHeader.GetChunkInfo()
+
+// Get total file size (for chunked uploads)
+totalSize := fileHeader.GetTotalSize()
+
+// Get current chunk size
+chunkSize := fileHeader.GetChunkSize()
+```
+
+## File Headers and Metadata
+
+The package supports several HTTP headers for enhanced functionality:
+
+- `Content-Range`: Standard HTTP range header for chunked uploads (e.g., "bytes 0-1023/2048")
+- `Content-Uid`: Unique identifier for file uploads (for deduplication and tracking)
+- `Content-Fingerprint`: Content-based hash for deduplication (e.g., "sha256:abc123")
+- `Content-Sync`: Enable synchronized uploads across multiple clients ("true"/"false")
+
+### Header Processing
+
+When processing uploads, headers can be extracted from both HTTP request headers and multipart file headers:
+
+```go
+// Extract headers from HTTP request and file headers
+header := attachment.GetHeader(requestHeader, fileHeader, fileSize)
+
+// The resulting FileHeader will contain merged headers from both sources
+uid := header.UID()
+fingerprint := header.Fingerprint()
+isSync := header.Sync()
+```
+
+## Upload Status Tracking
+
+Files have a status field that tracks the upload lifecycle:
+
+- `"uploading"`: File upload is in progress (for chunked uploads)
+- `"uploaded"`: File has been successfully uploaded
+- `"indexing"`: File is being processed for search indexing
+- `"indexed"`: File has been indexed and is fully processed
+- `"upload_failed"`: Upload failed due to an error
+- `"index_failed"`: Indexing failed but file is still accessible
+
+```go
+file, err := manager.Upload(ctx, fileHeader, reader, option)
+if err != nil {
+    return err
+}
+
+switch file.Status {
+case "uploading":
+    fmt.Println("Upload in progress...")
+case "uploaded":
+    fmt.Println("Upload completed successfully")
+case "upload_failed":
+    fmt.Println("Upload failed")
+}
+```
+
+#### `RegisterDefault(name string) (*Manager, error)`
+
+Registers a default attachment manager with sensible defaults for common file types.
--- a/neo/attachment/fileheader.go
+++ b/neo/attachment/fileheader.go
@ -10,6 +10,11 @@ func (fileheader *FileHeader) UID() string {
 	return fileheader.Header.Get("Content-Uid")
 }

+// Fingerprint is the fingerprint of the file, it is the fingerprint of the file
+func (fileheader *FileHeader) Fingerprint() string {
+	return fileheader.Header.Get("Content-Fingerprint")
+}
+
 // Range is the range of the file, it is the start and end of the file
 func (fileheader *FileHeader) Range() string {
 	return fileheader.Header.Get("Content-Range")
--- a/neo/attachment/manager.go
+++ b/neo/attachment/manager.go
@ -36,10 +36,10 @@ type UploadChunk struct {
 }

 // GetHeader gets the header from the file header and request header
-func GetHeader(requestHeader http.Header, fileHeader textproto.MIMEHeader) *FileHeader {
+func GetHeader(requestHeader http.Header, fileHeader textproto.MIMEHeader, size int64) *FileHeader {

 	// Convert the header to a FileHeader
-	header := &FileHeader{FileHeader: &multipart.FileHeader{Header: make(map[string][]string)}}
+	header := &FileHeader{FileHeader: &multipart.FileHeader{Header: make(map[string][]string), Size: size}}

 	for key, values := range fileHeader {
 		for _, value := range values {
@ -267,6 +267,10 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade
 			return nil, err
 		}

+		// Fix the file size, the file size is the sum of all chunks
+		file.Bytes = chunkIndex * int(chunkdata.Chunksize)
+		file.Status = "uploading"
+
 		// If this is the last chunk, merge all chunks
 		if fileheader.Complete() {
 			err = manager.storage.MergeChunks(ctx, file.ID, int(chunkdata.TotalChunks))
@ -284,6 +288,10 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade

 			// Remove the chunk data
 			uploadChunks.Delete(file.ID)
+
+			// Fix the file size
+			file.Bytes = int(chunkdata.Total)
+			file.Status = "uploaded"
 		}

 		return file, nil
@ -346,6 +354,7 @@ func (manager Manager) Upload(ctx context.Context, fileheader *FileHeader, reade

 	// Update the file ID
 	file.ID = id
+	file.Status = "uploaded"
 	return file, nil
 }

@ -487,6 +496,7 @@ func (manager Manager) makeFile(file *FileHeader, option UploadOption) (*File, e
 		ContentType: contentType,
 		Bytes:       int(file.Size),
 		CreatedAt:   int(time.Now().Unix()),
+		Status:      "uploading",
 	}, nil
 }

@ -516,7 +526,13 @@ func (manager Manager) allowed(contentType string, extension string) bool {

 // generateFileID generates a file ID with proper namespace
 func (manager Manager) generateFileID(file *FileHeader, extension string, option UploadOption) (string, error) {
-	filename := file.Filename
+
+	filename := file.Fingerprint()
+
+	// If the fingerprint is not set, use the filename
+	if filename == "" {
+		filename = file.Filename
+	}

 	// Use original filename if provided for better file identification
 	if option.OriginalFilename != "" {
--- a/neo/attachment/types.go
+++ b/neo/attachment/types.go
@ -8,14 +8,12 @@ import (

 // File the file
 type File struct {
-	ID          string   `json:"file_id"`
-	Bytes       int      `json:"bytes"`
-	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
+	ID          string `json:"file_id"`
+	Bytes       int    `json:"bytes"`
+	CreatedAt   int    `json:"created_at"`
+	Filename    string `json:"filename"`
+	ContentType string `json:"content_type"`
+	Status      string `json:"status"` // uploading, uploaded, indexing, indexed, upload_failed, index_failed
 }

 // FileResponse represents a file download response
--- a/neo/neo.go
+++ b/neo/neo.go
@ -50,12 +50,16 @@ func (neo *DSL) UserRoles(sid string) (interface{}, error) {
 }

 // UserOrGuestID get the user id or guest id from the session
-func (neo *DSL) UserOrGuestID(sid string) (interface{}, error) {
+func (neo *DSL) UserOrGuestID(sid string) (interface{}, bool, error) {
 	userID, err := neo.UserID(sid)
 	if err != nil {
-		return neo.GuestID(sid)
+		guestID, err := neo.GuestID(sid)
+		if err != nil {
+			return nil, false, err
+		}
+		return guestID, true, nil
 	}
-	return userID, nil
+	return userID, false, nil
 }

 // Download downloads a file
--- a/neo/types.go
+++ b/neo/types.go
@ -86,6 +86,14 @@ type Upload struct {
 	Knowledge *attachment.ManagerOption `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base upload setting, if not set use the chat upload setting.
 }

+// UploadOption the upload option
+type UploadOption struct {
+	attachment.UploadOption
+	Public       bool        `json:"public,omitempty" yaml:"public,omitempty, form:public"`                      // The public of the file, default is false
+	Scope        interface{} `json:"scope,omitempty" yaml:"scope,omitempty, form:scope"`                         // The scope of the file, default is private
+	CollectionID string      `json:"collection_id,omitempty" yaml:"collection_id,omitempty, form:collection_id"` // The collection id of the file, default is empty
+}
+
 // Knowledge base Settings
 // ===============================
 type Knowledge struct {