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"`