Skip to content

Commit

Permalink
Merge pull request #815 from trheyi/main
Browse files Browse the repository at this point in the history 
 Update Neo API with AWS Integration and Assistant Management
  • Loading branch information
@trheyi
trheyi authored Jan 3, 2025
2 parents 1254414 + 0d676b3 commit d7fecbf
Show file tree
Hide file tree
Showing 26 changed files with 2,163 additions and 882 deletions.
13 changes: 13 additions & 0 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@ go 1.23

require (
github.com/PuerkitoBio/goquery v1.10.1
github.com/aws/aws-sdk-go-v2 v1.32.7
github.com/aws/aws-sdk-go-v2/credentials v1.17.48
github.com/aws/aws-sdk-go-v2/service/s3 v1.71.1
github.com/blang/semver v3.5.1+incompatible
github.com/caarlos0/env/v6 v6.10.1
github.com/dchest/captcha v1.1.0
Expand Down Expand Up @@ -38,6 +41,15 @@ require (
filippo.io/edwards25519 v1.1.0 // indirect
github.com/TylerBrock/colorjson v0.0.0-20200706003622-8a50f05110d2 // indirect
github.com/andybalholm/cascadia v1.3.3 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.26 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.26 // indirect
github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.26 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.1 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.4.7 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.7 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.7 // indirect
github.com/aws/smithy-go v1.22.1 // indirect
github.com/blang/semver/v4 v4.0.0 // indirect
github.com/bytedance/sonic v1.12.6 // indirect
github.com/bytedance/sonic/loader v0.2.1 // indirect
Expand Down Expand Up @@ -90,6 +102,7 @@ require (
github.com/richardlehane/msoleps v1.0.4 // indirect
github.com/robfig/cron/v3 v3.0.1 // indirect
github.com/sirupsen/logrus v1.9.3 // indirect
github.com/spf13/cast v1.7.1 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/tcnksm/go-gitconfig v0.1.2 // indirect
github.com/tidwall/btree v1.7.0 // indirect
Expand Down
27 changes: 27 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,30 @@ github.com/TylerBrock/colorjson v0.0.0-20200706003622-8a50f05110d2 h1:ZBbLwSJqkH
github.com/TylerBrock/colorjson v0.0.0-20200706003622-8a50f05110d2/go.mod h1:VSw57q4QFiWDbRnjdX8Cb3Ow0SFncRw+bA/ofY6Q83w=
github.com/andybalholm/cascadia v1.3.3 h1:AG2YHrzJIm4BZ19iwJ/DAua6Btl3IwJX+VI4kktS1LM=
github.com/andybalholm/cascadia v1.3.3/go.mod h1:xNd9bqTn98Ln4DwST8/nG+H0yuB8Hmgu1YHNnWw0GeA=
github.com/aws/aws-sdk-go-v2 v1.32.7 h1:ky5o35oENWi0JYWUZkB7WYvVPP+bcRF5/Iq7JWSb5Rw=
github.com/aws/aws-sdk-go-v2 v1.32.7/go.mod h1:P5WJBrYqqbWVaOxgH0X/FYYD47/nooaPOZPlQdmiN2U=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7 h1:lL7IfaFzngfx0ZwUGOZdsFFnQ5uLvR0hWqqhyE7Q9M8=
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.7/go.mod h1:QraP0UcVlQJsmHfioCrveWOC1nbiWUl3ej08h4mXWoc=
github.com/aws/aws-sdk-go-v2/credentials v1.17.48 h1:IYdLD1qTJ0zanRavulofmqut4afs45mOWEI+MzZtTfQ=
github.com/aws/aws-sdk-go-v2/credentials v1.17.48/go.mod h1:tOscxHN3CGmuX9idQ3+qbkzrjVIx32lqDSU1/0d/qXs=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.26 h1:I/5wmGMffY4happ8NOCuIUEWGUvvFp5NSeQcXl9RHcI=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.26/go.mod h1:FR8f4turZtNy6baO0KJ5FJUmXH/cSkI9fOngs0yl6mA=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.26 h1:zXFLuEuMMUOvEARXFUVJdfqZ4bvvSgdGRq/ATcrQxzM=
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.26/go.mod h1:3o2Wpy0bogG1kyOPrgkXA8pgIfEEv0+m19O9D5+W8y8=
github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.26 h1:GeNJsIFHB+WW5ap2Tec4K6dzcVTsRbsT1Lra46Hv9ME=
github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.26/go.mod h1:zfgMpwHDXX2WGoG84xG2H+ZlPTkJUU4YUvx2svLQYWo=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.1 h1:iXtILhvDxB6kPvEXgsDhGaZCSC6LQET5ZHSdJozeI0Y=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.12.1/go.mod h1:9nu0fVANtYiAePIBh2/pFUSwtJ402hLnp854CNoDOeE=
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.4.7 h1:tB4tNw83KcajNAzaIMhkhVI2Nt8fAZd5A5ro113FEMY=
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.4.7/go.mod h1:lvpyBGkZ3tZ9iSsUIcC2EWp+0ywa7aK3BLT+FwZi+mQ=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.7 h1:8eUsivBQzZHqe/3FE+cqwfH+0p5Jo8PFM/QYQSmeZ+M=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.12.7/go.mod h1:kLPQvGUmxn/fqiCrDeohwG33bq2pQpGeY62yRO6Nrh0=
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.7 h1:Hi0KGbrnr57bEHWM0bJ1QcBzxLrL/k2DHvGYhb8+W1w=
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.7/go.mod h1:wKNgWgExdjjrm4qvfbTorkvocEstaoDl4WCvGfeCy9c=
github.com/aws/aws-sdk-go-v2/service/s3 v1.71.1 h1:aOVVZJgWbaH+EJYPvEgkNhCEbXXvH7+oML36oaPK3zE=
github.com/aws/aws-sdk-go-v2/service/s3 v1.71.1/go.mod h1:r+xl5yzMk9083rMR+sJ5TYj9Tihvf/l1oxzZXDgGj2Q=
github.com/aws/smithy-go v1.22.1 h1:/HPHZQ0g7f4eUeK6HKglFz8uwVfZKgoI25rb/J+dnro=
github.com/aws/smithy-go v1.22.1/go.mod h1:irrKGvNn1InZwb2d7fkIRNucdfwR8R+Ts3wxYa/cJHg=
github.com/blang/semver v3.5.1+incompatible h1:cQNTCjp13qL8KC3Nbxr/y2Bqb63oX6wdnnjpJbkM4JQ=
github.com/blang/semver v3.5.1+incompatible/go.mod h1:kRBLl5iJ+tD4TcOOxsy/0fnwebNt5EWlYSAyrTnjyyk=
github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM=
Expand Down Expand Up @@ -136,6 +160,7 @@ github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3x
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
Expand Down Expand Up @@ -195,6 +220,8 @@ github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncj
github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
github.com/spf13/cast v1.7.1 h1:cuNEagBQEHWN1FnbGEjCXL2szYEXqfJPbP2HNUaca9Y=
github.com/spf13/cast v1.7.1/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo=
github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
Expand Down
257 changes: 253 additions & 4 deletions neo/assistant/api.go
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,20 @@
package assistant

import (
"context"
"crypto/sha256"
"encoding/base64"
"fmt"
"io"
"mime/multipart"
"path/filepath"
"strings"
"time"

"github.com/yaoapp/gou/fs"
chatMessage "github.com/yaoapp/yao/neo/message"
)

// Get get the assistant by id
func Get(id string) (*Assistant, error) {
return LoadStore(id)
Expand All @@ -25,14 +40,248 @@ func GetByConnector(connector string, name string) (*Assistant, error) {
assistant, err := loadMap(data)
if err != nil {
return nil, err

}
loaded.Put(assistant)
return assistant, nil
}

// Init init the assistant
// Choose the connector and initialize the assistant
func (ast *Assistant) initialize() error {
// 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

// Chat implements the chat functionality
func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error {
if ast.openai == nil {
return fmt.Errorf("openai is not initialized")
}

requestMessages, err := ast.requestMessages(ctx, messages)
if err != nil {
return fmt.Errorf("request messages error: %s", err.Error())
}

_, ext := ast.openai.ChatCompletionsWith(ctx, requestMessages, option, cb)
if ext != nil {
return fmt.Errorf("openai chat completions with error: %s", ext.Message)
}

return nil
}

func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
newMessages := []map[string]interface{}{}

// With Prompts
if ast.Prompts != nil {
for _, prompt := range ast.Prompts {
message := map[string]interface{}{
"role": prompt.Role,
"content": prompt.Content,
}

name := ast.Name
if prompt.Name != "" {
name = prompt.Name
}

message["name"] = name
newMessages = append(newMessages, message)
}
}

length := len(messages)
for index, message := range messages {
role, ok := message["role"].(string)
if !ok {
return nil, fmt.Errorf("role must be string")
}

content, ok := message["content"].(string)
if !ok {
return nil, fmt.Errorf("content must be string")
}

newMessage := map[string]interface{}{
"role": role,
"content": content,
}

if name, ok := message["name"].(string); ok {
newMessage["name"] = name
}

// Special handling for user messages with JSON content last message
if role == "user" && index == length-1 {
content = strings.TrimSpace(content)
msg, err := chatMessage.NewString(content)
if err != nil {
return nil, fmt.Errorf("new string error: %s", err.Error())
}

newMessage["content"] = msg.Text
if msg.Attachments != nil {
content, err := ast.withAttachments(ctx, msg)
if err != nil {
return nil, fmt.Errorf("with attachments error: %s", err.Error())
}
newMessage["content"] = content
}
}

newMessages = append(newMessages, newMessage)
}
return newMessages, nil
}

func (ast *Assistant) withAttachments(ctx context.Context, msg *chatMessage.Message) ([]map[string]interface{}, error) {
contents := []map[string]interface{}{{"type": "text", "text": msg.Text}}
images := []string{}
for _, attachment := range msg.Attachments {
if strings.HasPrefix(attachment.ContentType, "image/") {
images = append(images, attachment.FileID)
}
}

if len(images) == 0 {
return contents, nil
}

for _, image := range images {
bytes64, err := ast.ReadBase64(ctx, image)
if err != nil {
return nil, fmt.Errorf("read base64 error: %s", err.Error())
}

contents = append(contents, map[string]interface{}{
"type": "image_url",
"image_url": map[string]string{
"url": fmt.Sprintf("data:image/jpeg;base64,%s", bytes64),
},
})
}

return contents, nil
}

// 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)
}

data, err := fs.Get("data")
if err != nil {
return nil, err
}

ext := filepath.Ext(file.Filename)
id, err := ast.id(file.Filename, ext)
if err != nil {
return nil, err
}

filename := id
_, err = data.Write(filename, reader, 0644)
if err != nil {
return nil, err
}

return &File{
ID: filename,
Filename: filename,
ContentType: contentType,
Bytes: int(file.Size),
CreatedAt: int(time.Now().Unix()),
}, 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
}

func (ast *Assistant) id(temp string, ext string) (string, error) {
date := time.Now().Format("20060102")
hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8]
return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), 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
}

// ReadBase64 implements base64 file reading functionality
func (ast *Assistant) ReadBase64(ctx context.Context, fileID string) (string, error) {
data, err := fs.Get("data")
if err != nil {
return "", fmt.Errorf("get filesystem error: %s", err.Error())
}

exists, err := data.Exists(fileID)
if err != nil {
return "", fmt.Errorf("check file error: %s", err.Error())
}
if !exists {
return "", fmt.Errorf("file %s not found", fileID)
}

content, err := data.ReadFile(fileID)
if err != nil {
return "", fmt.Errorf("read file error: %s", err.Error())
}

return base64.StdEncoding.EncodeToString(content), nil
}
2 changes: 1 addition & 1 deletion neo/assistant/assistant.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ func (ast *Assistant) Clone() *Assistant {
Mentionable: ast.Mentionable,
Automated: ast.Automated,
Script: ast.Script,
API: ast.API,
openai: ast.openai,
}

// Deep copy tags
Expand Down
Loading

0 comments on commit d7fecbf

Please sign in to comment.