Implement file download functionality and refactor existing methods in neo package
- Added handleDownload method to manage file download requests, including validation for session ID and file ID. - Enhanced permission checks to ensure users can only download files they are authorized to access. - Updated CORS headers to include Content-Disposition for better file handling in responses. - Removed deprecated download logic from the neo package to streamline code and improve maintainability. - Refactored existing methods to improve clarity and error handling during file operations.
This commit is contained in:
Max
parent
e6a0f37a66
commit
445bcdd7ea
4 changed files with 103 additions and 768 deletions
149
neo/api.go
149
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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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")
|
||||
})
|
||||
}
|
||||
28
neo/neo.go
28
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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue