opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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.

Browse source
This commit is contained in:
Max 2024-12-14 19:39:14 +08:00
parent c98a402e54
commit 773e050925
9 changed files with 305 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
neo/api.go
View file

@ -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

7
neo/assistant/base/base.go
View file

@ -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

86
neo/assistant/base/file.go Normal file
View file

@ -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
}

89
neo/assistant/openai/file.go
View file

@ -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() {}

7
neo/assistant/openai/openai.go
View file

@ -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

13
neo/assistant/types.go
View file

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

17
neo/hooks.go
View file

@ -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

105
neo/neo.go
View file

@ -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
}

9
neo/types.go
View file

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