From 773e05092512bd43052e38499e61d747e151443b Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 14 Dec 2024 19:39:14 +0800 Subject: [PATCH] Implement file upload functionality in Neo API, adding new endpoint for file uploads and enhancing assistant management. Introduce handleUpload method to process file uploads, including size and type validation. Update context to include upload information and refactor assistant selection logic for improved clarity. Enhance error handling and response structure for file upload operations, ensuring robust communication of success and failure states. --- neo/api.go | 26 +++++++- neo/assistant/base/base.go | 7 +-- neo/assistant/base/file.go | 86 +++++++++++++++++++++++++++ neo/assistant/openai/file.go | 89 ++++++++++++++++++++++++++-- neo/assistant/openai/openai.go | 7 +-- neo/assistant/types.go | 13 +++- neo/hooks.go | 17 +++--- neo/neo.go | 105 ++++++++++++++++++++++++--------- neo/types.go | 9 +++ 9 files changed, 305 insertions(+), 54 deletions(-) create mode 100644 neo/assistant/base/file.go diff --git a/neo/api.go b/neo/api.go index 36421a26..ab687947 100644 --- a/neo/api.go +++ b/neo/api.go @@ -27,6 +27,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.OPTIONS(path+"/status", neo.optionsHandler) router.OPTIONS(path+"/chats", neo.optionsHandler) router.OPTIONS(path+"/history", neo.optionsHandler) + router.OPTIONS(path+"/upload", neo.optionsHandler) // Register endpoints with middlewares router.GET(path, append(middlewares, neo.handleChat)...) @@ -34,7 +35,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.GET(path+"/status", append(middlewares, neo.handleStatus)...) router.GET(path+"/chats", append(middlewares, neo.handleChatList)...) router.GET(path+"/history", append(middlewares, neo.handleChatHistory)...) - + router.POST(path+"/upload", append(middlewares, neo.handleUpload)...) return nil } @@ -44,6 +45,29 @@ func (neo *DSL) handleStatus(c *gin.Context) { c.Done() } +// handleUpload handles the upload request +func (neo *DSL) handleUpload(c *gin.Context) { + sid := c.GetString("__sid") + if sid == "" { + sid = uuid.New().String() + } + + // Set the context + ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "") + defer cancel() + + // Upload the file + file, err := neo.Upload(ctx, c) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, file) + c.Done() +} + // handleChat handles the chat request func (neo *DSL) handleChat(c *gin.Context) { // Set headers for SSE diff --git a/neo/assistant/base/base.go b/neo/assistant/base/base.go index f6d1376a..bd608fe1 100644 --- a/neo/assistant/base/base.go +++ b/neo/assistant/base/base.go @@ -17,7 +17,7 @@ type Base struct { } // New create a new base assistant -func New(connector connector.Connector, prompts []assistant.Prompt, id ...string) (*Base, error) { +func New(connector connector.Connector, prompts []assistant.Prompt, id string) (*Base, error) { setting := connector.Setting() api, err := openai.NewOpenAI(setting) @@ -25,10 +25,7 @@ func New(connector connector.Connector, prompts []assistant.Prompt, id ...string return nil, err } - if len(id) > 0 { - return &Base{Connector: connector, ID: id[0], Prompts: prompts, openai: api}, nil - } - return &Base{Connector: connector, Prompts: prompts, openai: api}, nil + return &Base{Connector: connector, ID: id, Prompts: prompts, openai: api}, nil } // List list all assistants diff --git a/neo/assistant/base/file.go b/neo/assistant/base/file.go new file mode 100644 index 00000000..487182de --- /dev/null +++ b/neo/assistant/base/file.go @@ -0,0 +1,86 @@ +package base + +import ( + "context" + "crypto/sha256" + "fmt" + "io" + "mime/multipart" + "path/filepath" + "strings" + "time" + + "github.com/yaoapp/gou/fs" + "github.com/yaoapp/yao/neo/assistant" +) + +// AllowedFileTypes the allowed file types +var AllowedFileTypes = map[string]string{ + "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 the file +func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.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) + if err != nil { + return nil, err + } + + filename := fmt.Sprintf("%s%s", id, ext) + _, err = data.Write(filename, reader, 0644) + if err != nil { + return nil, err + } + + return &assistant.File{ + ID: strings.ReplaceAll(id, "/", "_"), + Filename: filename, + ContentType: contentType, + Bytes: int(file.Size), + CreatedAt: int(time.Now().Unix()), + }, nil +} + +func (ast *Base) id(temp string) (string, error) { + date := time.Now().Format("20060102") + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8] + return fmt.Sprintf("/__assistants/%s/%s/%s", ast.ID, date, hash), nil +} + +func (ast *Base) allowed(contentType string) bool { + if _, ok := AllowedFileTypes[contentType]; ok { + return true + } + // text/* // image/* // audio/* // video/* + 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/openai/file.go b/neo/assistant/openai/file.go index 3b040d39..9c52801d 100644 --- a/neo/assistant/openai/file.go +++ b/neo/assistant/openai/file.go @@ -1,16 +1,93 @@ package openai -// File the file struct -type File struct { - ID string `json:"file_id"` +import ( + "context" + "crypto/sha256" + "fmt" + "io" + "mime/multipart" + "path/filepath" + "strings" + "time" + + "github.com/yaoapp/gou/fs" + "github.com/yaoapp/yao/neo/assistant" +) + +// AllowedFileTypes the allowed file types +var AllowedFileTypes = map[string]string{ + "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 the file +func (ast *OpenAI) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.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) + if err != nil { + return nil, err + } + + filename := fmt.Sprintf("%s%s", id, ext) + _, err = data.Write(filename, reader, 0644) + if err != nil { + return nil, err + } + + return &assistant.File{ + ID: strings.ReplaceAll(id, "/", "_"), + Filename: filename, + ContentType: contentType, + Bytes: int(file.Size), + CreatedAt: int(time.Now().Unix()), + }, nil +} + +func (ast *OpenAI) id(temp string) (string, error) { + date := time.Now().Format("20060102") + hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8] + return fmt.Sprintf("/__assistants/%s/%s/%s", ast.ID, date, hash), nil +} + +func (ast *OpenAI) allowed(contentType string) bool { + if _, ok := AllowedFileTypes[contentType]; ok { + return true + } + // text/* // image/* // audio/* // video/* + if strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "audio/") || strings.HasPrefix(contentType, "video/") { + return true + } + return false } // FileLists list all files func (ast *OpenAI) FileLists() {} -// Upload upload a file to an assistant -func (ast *OpenAI) Upload() {} - // FileDelete delete a file func (ast *OpenAI) FileDelete() {} diff --git a/neo/assistant/openai/openai.go b/neo/assistant/openai/openai.go index ee10d12c..38f5135c 100644 --- a/neo/assistant/openai/openai.go +++ b/neo/assistant/openai/openai.go @@ -16,7 +16,7 @@ type OpenAI struct { } // New create a new openai assistant -func New(connector connector.Connector, id ...string) (*OpenAI, error) { +func New(connector connector.Connector, id string) (*OpenAI, error) { setting := connector.Setting() openai, err := api.NewOpenAI(setting) @@ -24,10 +24,7 @@ func New(connector connector.Connector, id ...string) (*OpenAI, error) { return nil, err } - if len(id) > 0 { - return &OpenAI{ID: id[0], Connector: connector, openai: openai}, nil - } - return &OpenAI{Connector: connector, openai: openai}, nil + return &OpenAI{ID: id, Connector: connector, openai: openai}, nil } // Current set the current assistant diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 33566c05..3e338cb0 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -2,12 +2,14 @@ package assistant import ( "context" + "io" + "mime/multipart" ) // 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 - List(ctx context.Context, param QueryParam) ([]Assistant, error) + Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error) } // Prompt a prompt @@ -35,3 +37,12 @@ type Assistant struct { Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts API API `json:"-" yaml:"-"` // Assistant API } + +// 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"` +} diff --git a/neo/hooks.go b/neo/hooks.go index c189c74f..7066cf05 100644 --- a/neo/hooks.go +++ b/neo/hooks.go @@ -12,8 +12,16 @@ import ( // HookCreate create the assistant func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) { + + // Default assistant + assistantID := neo.Use + if ctx.AssistantID != "" { + assistantID = ctx.AssistantID + } + + // Empty hook if neo.Create == "" { - return CreateResponse{AssistantID: neo.Use, ChatID: ctx.ChatID}, nil + return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil } // Create a context with 10 second timeout @@ -42,14 +50,10 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi return v, nil case map[string]interface{}: - assistantID := "" if id, ok := v["assistant_id"].(string); ok { assistantID = id } - if assistantID == "" && neo.Use != "" { - assistantID = neo.Use - } chatID := "" if id, ok := v["chat_id"].(string); ok { chatID = id @@ -62,8 +66,7 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi return CreateResponse{AssistantID: assistantID, ChatID: chatID}, nil } - // Default assistant - return CreateResponse{AssistantID: neo.Use, ChatID: ctx.ChatID}, nil + return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil } // HookAssistants query the assistant list from the assistant list hook diff --git a/neo/neo.go b/neo/neo.go index e1021dc0..6a87be8a 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -2,6 +2,7 @@ package neo import ( "fmt" + "os" "strings" "sync" @@ -37,20 +38,62 @@ func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { } // Select Assistant - ast := neo.Assistant - if res.AssistantID != "" { - ast, err = neo.newAssistant(res.AssistantID) - if err != nil { - msg := message.New().Error(err).Done() - msg.Write(c.Writer) - return err - } + ast, err := neo.selectAssistant(res.AssistantID) + if err != nil { + return err } // Chat with AI return neo.chat(ast, ctx, messages, c) } +// Upload upload a file +func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) { + // Get the file + tmpfile, err := c.FormFile("file") + if err != nil { + return nil, err + } + + reader, err := tmpfile.Open() + if err != nil { + return nil, err + } + defer func() { + reader.Close() + os.Remove(tmpfile.Filename) + }() + + // Get option from form data option_xxx + option := map[string]interface{}{} + for key := range c.Request.Form { + if strings.HasPrefix(key, "option_") { + option[strings.TrimPrefix(key, "option_")] = c.PostForm(key) + } + } + + // Get file info + ctx.Upload = &FileUpload{ + Bytes: int(tmpfile.Size), + Name: tmpfile.Filename, + ContentType: tmpfile.Header.Get("Content-Type"), + Option: option, + } + + res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c) + if err != nil { + return nil, err + } + + // Select Assistant + ast, err := neo.selectAssistant(res.AssistantID) + if err != nil { + return nil, err + } + + return ast.Upload(ctx, tmpfile, reader, option) +} + // chat chat with AI func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error { @@ -143,6 +186,19 @@ func (neo *DSL) updateAssistantList(list []assistant.Assistant) { } } +// selectAssistant select the assistant +func (neo *DSL) selectAssistant(assistantID string) (assistant.API, error) { + ast := neo.Assistant + if assistantID != "" { + ast, err := neo.newAssistant(assistantID) + if err != nil { + return nil, err + } + return ast, nil + } + return ast, nil +} + // newAssistant create a new assistant func (neo *DSL) newAssistant(id string) (assistant.API, error) { // Try to find assistant in AssistantList first @@ -182,7 +238,7 @@ func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { } if conn.Is(connector.OPENAI) { - api, err := openai.New(conn, neo.Use) + api, err := openai.New(conn, id) if err != nil { return nil, fmt.Errorf("Create openai assistant error: %s", err.Error()) } @@ -190,7 +246,7 @@ func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { } // Base on the assistant list hook - api, err := base.New(conn, neo.Prompts, neo.Use) + api, err := base.New(conn, neo.Prompts, id) if err != nil { return nil, fmt.Errorf("Create base assistant error: %s", err.Error()) } @@ -226,7 +282,7 @@ func (neo *DSL) newMoapiAssistant(id string) (assistant.API, error) { return nil, fmt.Errorf("Create moapi assistant error: %s", err.Error()) } - api, err := openai.New(conn, neo.Use) + api, err := openai.New(conn, strings.ReplaceAll(id, ":", "_")) if err != nil { return nil, fmt.Errorf("Create openai assistant error: %s", err.Error()) } @@ -241,31 +297,22 @@ func (neo *DSL) createDefaultAssistant() (assistant.API, error) { return neo.newAssistant(neo.Connector) } -// prompts get the prompts -func (neo *DSL) prompts() []map[string]interface{} { - prompts := []map[string]interface{}{} - for _, prompt := range neo.Prompts { - message := map[string]interface{}{"role": prompt.Role, "content": prompt.Content} - if prompt.Name != "" { - message["name"] = prompt.Name - } - prompts = append(prompts, message) - } - - return prompts -} - // chatMessages get the chat messages -func (neo *DSL) chatMessages(ctx Context, content string) ([]map[string]interface{}, error) { +func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) { history, err := neo.Conversation.GetHistory(ctx.Sid, ctx.ChatID) if err != nil { return nil, err } - messages := append([]map[string]interface{}{}, neo.prompts()...) - messages = append(messages, history...) - messages = append(messages, map[string]interface{}{"role": "user", "content": content, "name": ctx.Sid}) + 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 } diff --git a/neo/types.go b/neo/types.go index c0f9694a..4e7c7c22 100644 --- a/neo/types.go +++ b/neo/types.go @@ -42,6 +42,7 @@ type Context struct { 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:"-"` } @@ -51,6 +52,14 @@ type Field struct { Bind string `json:"bind,omitempty"` } +// FileUpload the file upload info +type FileUpload struct { + Bytes int `json:"bytes,omitempty"` // If upload file, the file bytes + Name string `json:"name,omitempty"` // If upload + ContentType string `json:"content_type,omitempty"` // If upload file, the file content type + Option map[string]interface{} `json:"option,omitempty"` // If upload file, the upload option +} + // CreateResponse the response of the create hook type CreateResponse struct { AssistantID string `json:"assistant_id,omitempty"`