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

Merge pull request #815 from trheyi/main

Browse source 
Update Neo API with AWS Integration and Assistant Management
This commit is contained in:
Max 2025-01-03 17:29:12 +08:00 • committed by GitHub
commit d7fecbf709
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 2163 additions and 882 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
go.mod
View file

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

27
go.sum
View file

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

257
neo/assistant/api.go
View file

@ -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)
@ -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
neo/assistant/assistant.go
View file

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

72
neo/assistant/load.go
View file

@ -8,10 +8,13 @@ import (
"time"
jsoniter "github.com/json-iterator/go"
"github.com/spf13/cast"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/gou/rag/driver"
v8 "github.com/yaoapp/gou/runtime/v8"
"github.com/yaoapp/yao/neo/store"
neovision "github.com/yaoapp/yao/neo/vision"
"github.com/yaoapp/yao/openai"
"github.com/yaoapp/yao/share"
"gopkg.in/yaml.v3"
)
@ -20,6 +23,8 @@ import (
var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil
var rag *RAG = nil
var vision *neovision.Vision = nil
var defaultConnector string = "" // default connector
// LoadBuiltIn load the built-in assistants
func LoadBuiltIn() error {
@ -62,17 +67,39 @@ func LoadBuiltIn() error {
assistant.Readonly = true
assistant.BuiltIn = true
assistant.Sort = sort
if assistant.Sort == 0 {
assistant.Sort = sort
}
if assistant.Tags == nil {
assistant.Tags = []string{"Built-in"}
}
// Check if the assistant has Built-in tag
hasBuiltIn := false
for _, tag := range assistant.Tags {
if tag == "Built-in" {
hasBuiltIn = true
break
}
}
// add Built-in tag if not exists
if !hasBuiltIn {
assistant.Tags = append(assistant.Tags, "Built-in")
}
// Save the assistant
err = assistant.Save()
if err != nil {
return err
}
// Initialize the assistant
err = assistant.initialize()
if err != nil {
return err
}
sort++
loaded.Put(assistant)
@ -86,6 +113,16 @@ func SetStorage(s store.Store) {
storage = s
}
// SetVision set the vision
func SetVision(v *neovision.Vision) {
vision = v
}
// SetConnector set the connector
func SetConnector(c string) {
defaultConnector = c
}
// SetRAG set the RAG engine
// e: the RAG engine
// u: the RAG file uploader
@ -115,6 +152,11 @@ func ClearCache() {
// LoadStore create a new assistant from store
func LoadStore(id string) (*Assistant, error) {
if id == "" {
return nil, fmt.Errorf("assistant_id is required")
}
assistant, exists := loaded.Get(id)
if exists {
return assistant, nil
@ -259,8 +301,8 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
}
// sort
if v, ok := data["sort"].(int); ok {
assistant.Sort = v
if v, has := data["sort"]; has {
assistant.Sort = cast.ToInt(v)
}
// path
@ -331,6 +373,12 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.UpdatedAt = ts
}
// Initialize the assistant
err := assistant.initialize()
if err != nil {
return nil, err
}
return assistant, nil
}
@ -398,3 +446,21 @@ func loadScriptSource(source string, file string) (*v8.Script, error) {
}
return script, nil
}
// Init init the assistant
// Choose the connector and initialize the assistant
func (ast *Assistant) initialize() error {
conn := defaultConnector
if ast.Connector != "" {
conn = ast.Connector
}
ast.Connector = conn
api, err := openai.New(conn)
if err != nil {
return err
}
ast.openai = api
return nil
}

21
neo/assistant/local/chat.go
View file

@ -1,21 +0,0 @@
package local
import (
"context"
"fmt"
)
// Chat the chat
func (ast *Local) 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("api is not initialized")
}
_, ext := ast.openai.ChatCompletionsWith(ctx, messages, option, cb)
if ext != nil {
return fmt.Errorf("openai chat completions with error: %s", ext.Message)
}
return nil
}

162
neo/assistant/local/file.go
View file

@ -1,162 +0,0 @@
package local
import (
"context"
"crypto/sha256"
"encoding/base64"
"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/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 the file
func (ast *Local) 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, ext)
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: filename,
Filename: filename,
ContentType: contentType,
Bytes: int(file.Size),
CreatedAt: int(time.Now().Unix()),
}, nil
}
func (ast *Local) 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
}
func (ast *Local) 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
}
// Download downloads a file
func (ast *Local) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) {
// Get the data filesystem
data, err := fs.Get("data")
if err != nil {
return nil, fmt.Errorf("get filesystem error: %s", err.Error())
}
// Check if file exists
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)
}
// Open the file
reader, err := data.ReadCloser(fileID)
if err != nil {
return nil, err
}
// Get content type and extension
ext := filepath.Ext(fileID)
// Get content type from mime type
contentType := "application/octet-stream"
if v, err := data.MimeType(fileID); err == nil {
contentType = v
}
for mimeType, extension := range AllowedFileTypes {
if "."+extension == ext {
contentType = mimeType
break
}
}
return &assistant.FileResponse{
Reader: reader,
ContentType: contentType,
Extension: ext,
}, nil
}
// ReadBase64 reads a file and returns its base64 encoded content
func (ast *Local) ReadBase64(ctx context.Context, fileID string) (string, error) {
// Get the data filesystem
data, err := fs.Get("data")
if err != nil {
return "", fmt.Errorf("get filesystem error: %s", err.Error())
}
// Check if file exists
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)
}
// Read file content
content, err := data.ReadFile(fileID)
if err != nil {
return "", fmt.Errorf("read file error: %s", err.Error())
}
// Encode to base64
return base64.StdEncoding.EncodeToString(content), nil
}

34
neo/assistant/local/local.go
View file

@ -1,34 +0,0 @@
package local
import (
"context"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/openai"
)
// Local the local assistant
type Local struct {
ID string `json:"assistant_id"`
Prompts []assistant.Prompt `json:"prompts,omitempty"`
Connector connector.Connector `json:"-" yaml:"-"`
openai *openai.OpenAI
}
// New create a new local assistant
func New(connector connector.Connector, prompts []assistant.Prompt, id string) (*Local, error) {
setting := connector.Setting()
api, err := openai.NewOpenAI(setting)
if err != nil {
return nil, err
}
return &Local{Connector: connector, ID: id, Prompts: prompts, openai: api}, nil
}
// List list all assistants
func (ast *Local) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) {
return nil, nil
}

118
neo/assistant/openai/chat.go
View file

@ -1,118 +0,0 @@
package openai
import (
"context"
"fmt"
"strings"
chatMessage "github.com/yaoapp/yao/neo/message"
)
// Chat the chat struct
type Chat struct {
ID string `json:"chat_id"`
ThreadID string `json:"thread_id"`
}
// NewChat create a new chat
func (ast *OpenAI) NewChat() {}
// Chat the chat
func (ast *OpenAI) 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 *OpenAI) requestMessages(ctx context.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
newMessages := []map[string]interface{}{}
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,
}
// Handle name if present
if name, ok := message["name"].(string); ok {
newMessage["name"] = name
}
newMessage["content"] = content
// 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 *OpenAI) 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
}

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

@ -1,167 +0,0 @@
package openai
import (
"context"
"crypto/sha256"
"encoding/base64"
"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/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 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, ext)
if err != nil {
return nil, err
}
filename := id
_, err = data.Write(filename, reader, 0644)
if err != nil {
return nil, err
}
return &assistant.File{
ID: filename,
Filename: filename,
ContentType: contentType,
Bytes: int(file.Size),
CreatedAt: int(time.Now().Unix()),
}, nil
}
func (ast *OpenAI) 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
}
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() {}
// FileDelete delete a file
func (ast *OpenAI) FileDelete() {}
// FileContent get the content of a file
func (ast *OpenAI) FileContent() {}
// FileInfo get the information of a file
func (ast *OpenAI) FileInfo() {}
// Download downloads a file
func (ast *OpenAI) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) {
// Get the data filesystem
data, err := fs.Get("data")
if err != nil {
return nil, fmt.Errorf("get filesystem error: %s", err.Error())
}
// Check if file exists
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)
}
// Open the file
reader, err := data.ReadCloser(fileID)
if err != nil {
return nil, err
}
// Get content type and extension
ext := filepath.Ext(fileID)
// Get content type from mime type
contentType := "application/octet-stream"
if v, err := data.MimeType(fileID); err == nil {
contentType = v
}
return &assistant.FileResponse{
Reader: reader,
ContentType: contentType,
Extension: ext,
}, nil
}
// ReadBase64 reads a file and returns its base64 encoded content
func (ast *OpenAI) ReadBase64(ctx context.Context, fileID string) (string, error) {
// Get the data filesystem
data, err := fs.Get("data")
if err != nil {
return "", fmt.Errorf("get filesystem error: %s", err.Error())
}
// Check if file exists
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)
}
// Read file content
content, err := data.ReadFile(fileID)
if err != nil {
return "", fmt.Errorf("read file error: %s", err.Error())
}
// Encode to base64
return base64.StdEncoding.EncodeToString(content), nil
}

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

@ -1,51 +0,0 @@
package openai
import (
"context"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/yao/neo/assistant"
api "github.com/yaoapp/yao/openai"
)
// OpenAI the openai assistant
type OpenAI struct {
ID string `json:"assistant_id"` // the assistant id
Connector connector.Connector `json:"-" yaml:"-"`
openai *api.OpenAI
}
// New create a new openai assistant
func New(connector connector.Connector, id string) (*OpenAI, error) {
setting := connector.Setting()
openai, err := api.NewOpenAI(setting)
if err != nil {
return nil, err
}
return &OpenAI{ID: id, Connector: connector, openai: openai}, nil
}
// Current set the current assistant
func (ast *OpenAI) Current(id string) *OpenAI {
ast.ID = id
return ast
}
// List list all assistants
func (ast *OpenAI) List(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) {
return nil, nil
}
// Create create a new assistant
func (ast *OpenAI) Create() {}
// Delete delete an assistant
func (ast *OpenAI) Delete() {}
// Update update an assistant
func (ast *OpenAI) Update() {}
// Get get an assistant
func (ast *OpenAI) Get() {}

21
neo/assistant/openai/thread.go
View file

@ -1,21 +0,0 @@
package openai
// Thread the thread struct
type Thread struct {
ID string `json:"thread_id"`
}
// ThreadList list all threads
func (ast *OpenAI) ThreadList() {}
// ThreadCreate create a new thread
func (ast *OpenAI) ThreadCreate() {}
// ThreadGet get a thread
func (ast *OpenAI) ThreadGet(id string) {}
// ThreadDelete delete a thread
func (ast *OpenAI) ThreadDelete() {}
// ThreadUpdate update a thread
func (ast *OpenAI) ThreadUpdate() {}

3
neo/assistant/types.go
View file

@ -7,6 +7,7 @@ import (
"github.com/yaoapp/gou/rag/driver"
v8 "github.com/yaoapp/gou/runtime/v8"
api "github.com/yaoapp/yao/openai"
)
// API the assistant API interface
@ -64,9 +65,9 @@ type Assistant struct {
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows
Script *v8.Script `json:"-" yaml:"-"` // Assistant Script
API API `json:"-" yaml:"-"` // Assistant API
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
openai *api.OpenAI // OpenAI API
}
// File the file

103
neo/hooks.go
View file

@ -7,7 +7,6 @@ import (
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/neo/assistant"
)
// HookCreate create the assistant
@ -69,51 +68,6 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi
return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil
}
// HookAssistants query the assistant list from the assistant list hook
func (neo *DSL) HookAssistants(ctx context.Context, param assistant.QueryParam) ([]assistant.Assistant, error) {
if neo.AssistantListHook == "" {
return nil, nil
}
// Create a context with 10 second timeout
timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
p, err := process.Of(neo.AssistantListHook, param)
if err != nil {
return nil, err
}
err = p.WithContext(timeoutCtx).Execute()
if err != nil {
return nil, err
}
defer p.Release()
// Check if context was canceled
if timeoutCtx.Err() != nil {
return nil, timeoutCtx.Err()
}
value := p.Value()
if value == nil {
return nil, nil
}
var list []assistant.Assistant
bytes, err := jsoniter.Marshal(value)
if err != nil {
return nil, err
}
err = jsoniter.Unmarshal(bytes, &list)
if err != nil {
return nil, err
}
return list, nil
}
// HookPrepare executes the prepare hook before AI is called
func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
if neo.Prepare == "" {
@ -194,60 +148,3 @@ func (neo *DSL) HookWrite(ctx Context, messages []map[string]interface{}, respon
return result, nil
}
// HookMention query the mention list
func (neo *DSL) HookMention(ctx context.Context, keywords string) ([]Mention, error) {
// Default Get the assistant list
if neo.MentionHook == "" {
var mentions []Mention
assistants := neo.GetAssistants()
for _, assistant := range assistants {
mentions = append(mentions, Mention{
ID: assistant.ID,
Name: assistant.Name,
Type: "assistant",
})
}
return mentions, nil
}
// Create a context with 10 second timeout
timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
p, err := process.Of(neo.MentionHook, keywords)
if err != nil {
return nil, err
}
err = p.WithContext(timeoutCtx).Execute()
if err != nil {
return nil, err
}
defer p.Release()
// Check if context was canceled
if timeoutCtx.Err() != nil {
return nil, timeoutCtx.Err()
}
value := p.Value()
if value == nil {
return nil, nil
}
var list []Mention
bytes, err := jsoniter.Marshal(value)
if err != nil {
return nil, err
}
err = jsoniter.Unmarshal(bytes, &list)
if err != nil {
return nil, err
}
return list, nil
}

36
neo/load.go
View file

@ -12,6 +12,8 @@ import (
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/rag"
"github.com/yaoapp/yao/neo/store"
"github.com/yaoapp/yao/neo/vision"
"github.com/yaoapp/yao/neo/vision/driver"
)
// Neo the neo AI assistant
@ -56,6 +58,9 @@ func Load(cfg config.Config) error {
// Initialize RAG
Neo.initRAG()
// Initialize Vision
Neo.initVision()
// Initialize Assistant
err = Neo.initAssistant()
if err != nil {
@ -111,6 +116,27 @@ func (neo *DSL) initStore() error {
return fmt.Errorf("%s store connector %s not support", neo.ID, neo.StoreSetting.Connector)
}
// initVision initialize the Vision instance
func (neo *DSL) initVision() {
if neo.VisionSetting.Storage.Driver == "" {
return
}
cfg := &driver.Config{
Storage: neo.VisionSetting.Storage,
Model: neo.VisionSetting.Model,
}
instance, err := vision.New(cfg)
if err != nil {
color.Red("[Neo] Failed to initialize Vision: %v", err)
log.Error("[Neo] Failed to initialize Vision: %v", err)
return
}
neo.Vision = instance
}
// initAssistant initialize the assistant
func (neo *DSL) initAssistant() error {
@ -129,6 +155,14 @@ func (neo *DSL) initAssistant() error {
)
}
// Assistant Vision
if Neo.Vision != nil {
assistant.SetVision(Neo.Vision)
}
// Default Connector
assistant.SetConnector(Neo.Connector)
// Load Built-in Assistants
err := assistant.LoadBuiltIn()
if err != nil {
@ -141,7 +175,7 @@ func (neo *DSL) initAssistant() error {
return err
}
Neo.Assistant = defaultAssistant.API
Neo.Assistant = defaultAssistant
return nil
}

184
neo/neo.go
View file

@ -1,21 +1,15 @@
package neo
import (
"context"
"fmt"
"os"
"strings"
"sync"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/assistant/local"
"github.com/yaoapp/yao/neo/assistant/openai"
"github.com/yaoapp/yao/neo/message"
"github.com/yaoapp/yao/neo/store"
"github.com/yaoapp/yao/share"
)
// Lock the assistant list
@ -39,23 +33,20 @@ func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error {
}
// Select Assistant
ast, err := neo.selectAssistant(res.AssistantID)
ast, err := neo.Select(res.AssistantID)
if err != nil {
return err
}
// Chat with AI
return neo.chat(ast, ctx, messages, c)
}
// GetAssistants returns the list of assistants
func (neo *DSL) GetAssistants() []assistant.Assistant {
return neo.AssistantList
}
// GetMentions returns the mention list
func (neo *DSL) GetMentions(keywords string) ([]Mention, error) {
return neo.HookMention(context.Background(), keywords)
// Select select an assistant
func (neo *DSL) Select(id string) (assistant.API, error) {
if id == "" {
return neo.Assistant, nil
}
return assistant.Get(id)
}
// GeneratePrompts generate prompts for the AI assistant
@ -110,7 +101,7 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy
}
// Select Assistant
ast, err := neo.selectAssistant(res.AssistantID)
ast, err := neo.Select(res.AssistantID)
if err != nil {
return "", err
}
@ -234,7 +225,7 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
}
// Select Assistant
ast, err := neo.selectAssistant(res.AssistantID)
ast, err := neo.Select(res.AssistantID)
if err != nil {
return nil, err
}
@ -257,7 +248,7 @@ func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse,
}
// Select Assistant
ast, err := neo.selectAssistant(res.AssistantID)
ast, err := neo.Select(res.AssistantID)
if err != nil {
return nil, err
}
@ -346,130 +337,6 @@ func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]inter
}
}
// updateAssistantList update the assistant list
func (neo *DSL) updateAssistantList(list []assistant.Assistant) {
lock.Lock()
defer lock.Unlock()
neo.AssistantList = list
neo.AssistantMaps = make(map[string]assistant.Assistant)
if list != nil {
for _, assistant := range list {
neo.AssistantMaps[assistant.ID] = 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
if id != "" && neo.AssistantMaps != nil {
if ast, ok := neo.AssistantMaps[id]; ok {
if ast.API != nil {
return ast.API, nil
}
api, err := neo.newAssistantByConfig(&ast)
if err != nil {
return nil, err
}
ast.API = api
return api, nil
}
}
return neo.newAssistantByConnector(id)
}
// newAssistantByConfig create a new assistant from assistant configuration
func (neo *DSL) newAssistantByConfig(ast *assistant.Assistant) (assistant.API, error) {
return neo.newAssistantByConnector(ast.Connector)
}
// newAssistantByConnector create a new assistant from connector id
func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) {
// Moapi connector
if id == "" || strings.HasPrefix(id, "moapi") {
return neo.newMoapiAssistant(id)
}
// Other connector
conn, err := connector.Select(id)
if err != nil {
return nil, fmt.Errorf("Neo assistant connector %s not support", id)
}
if conn.Is(connector.OPENAI) {
api, err := openai.New(conn, id)
if err != nil {
return nil, fmt.Errorf("Create openai assistant error: %s", err.Error())
}
return api, nil
}
// Base on the assistant list hook
api, err := local.New(conn, neo.Prompts, id)
if err != nil {
return nil, fmt.Errorf("Create local assistant error: %s", err.Error())
}
return api, nil
}
// newMoapiAssistant creates a new moapi assistant
func (neo *DSL) newMoapiAssistant(id string) (assistant.API, error) {
model := "gpt-3.5-turbo"
if strings.HasPrefix(id, "moapi:") {
model = strings.TrimPrefix(id, "moapi:")
}
// Get the moapi setting
url := share.MoapiHosts[0]
if share.App.Moapi.Mirrors != nil {
url = share.App.Moapi.Mirrors[0]
}
key := share.App.Moapi.Secret
organization := share.App.Moapi.Organization
if !strings.HasPrefix(url, "http") {
url = "https://" + url
}
// Check the moapi secret
if key == "" {
return nil, fmt.Errorf("The moapi secret is empty")
}
conn, err := connector.New(`moapi`, `__yao.moapi`, []byte(`{"name":"Moapi", "options":{"model": "`+model+`", "key": "`+key+`", "organization": "`+organization+`", "host": "`+url+`"}}`))
if err != nil {
return nil, fmt.Errorf("Create moapi assistant error: %s", err.Error())
}
api, err := openai.New(conn, strings.ReplaceAll(id, ":", "_"))
if err != nil {
return nil, fmt.Errorf("Create openai assistant error: %s", err.Error())
}
return api, nil
}
// createDefaultAssistant create a default assistant
func (neo *DSL) createDefaultAssistant() (assistant.API, error) {
if neo.Use != "" {
return neo.newAssistant(neo.Use)
}
return neo.newAssistant(neo.Connector)
}
// chatMessages get the chat messages
func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) {
@ -509,37 +376,6 @@ func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages
}
}
// createStore create a new store
func (neo *DSL) createStore() error {
var err error
if neo.StoreSetting.Connector == "default" || neo.StoreSetting.Connector == "" {
neo.Store, err = store.NewXun(neo.StoreSetting)
return err
}
// other connector
conn, err := connector.Select(neo.StoreSetting.Connector)
if err != nil {
return err
}
if conn.Is(connector.DATABASE) {
neo.Store, err = store.NewXun(neo.StoreSetting)
return err
} else if conn.Is(connector.REDIS) {
neo.Store = store.NewRedis()
return nil
} else if conn.Is(connector.MONGO) {
neo.Store = store.NewMongo()
return nil
}
return fmt.Errorf("%s store connector %s not support", neo.ID, neo.StoreSetting.Connector)
}
// sendMessage sends a message to the client
func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error {
if msg, ok := data.(map[string]interface{}); ok {

48
neo/types.go
View file

@ -7,31 +7,37 @@ import (
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/rag"
"github.com/yaoapp/yao/neo/store"
"github.com/yaoapp/yao/neo/vision"
"github.com/yaoapp/yao/neo/vision/driver"
)
// DSL AI assistant
type DSL struct {
ID string `json:"-" yaml:"-"`
Name string `json:"name,omitempty" yaml:"name,omitempty"`
Use string `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default
Guard string `json:"guard,omitempty" yaml:"guard,omitempty"`
Connector string `json:"connector" yaml:"connector"`
StoreSetting store.Setting `json:"store" yaml:"store"`
RAGSetting rag.Setting `json:"rag" yaml:"rag"`
Option map[string]interface{} `json:"option" yaml:"option"`
Prepare string `json:"prepare,omitempty" yaml:"prepare,omitempty"`
Create string `json:"create,omitempty" yaml:"create,omitempty"`
Write string `json:"write,omitempty" yaml:"write,omitempty"`
AssistantListHook string `json:"assistants,omitempty" yaml:"assistants,omitempty"` // Get the assistant list from the hook
MentionHook string `json:"mentions,omitempty"` // Get the mention list from the hook
Prompts []assistant.Prompt `json:"prompts,omitempty" yaml:"prompts,omitempty"`
Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"`
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
Store store.Store `json:"-" yaml:"-"`
RAG *rag.RAG `json:"-" yaml:"-"`
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
AssistantList []assistant.Assistant `json:"-" yaml:"-"`
AssistantMaps map[string]assistant.Assistant `json:"-" yaml:"-"`
ID string `json:"-" yaml:"-"`
Name string `json:"name,omitempty" yaml:"name,omitempty"`
Use string `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default
Guard string `json:"guard,omitempty" yaml:"guard,omitempty"`
Connector string `json:"connector" yaml:"connector"`
StoreSetting store.Setting `json:"store" yaml:"store"`
RAGSetting rag.Setting `json:"rag" yaml:"rag"`
VisionSetting VisionSetting `json:"vision" yaml:"vision"`
Option map[string]interface{} `json:"option" yaml:"option"`
Prepare string `json:"prepare,omitempty" yaml:"prepare,omitempty"`
Create string `json:"create,omitempty" yaml:"create,omitempty"`
Write string `json:"write,omitempty" yaml:"write,omitempty"`
Prompts []assistant.Prompt `json:"prompts,omitempty" yaml:"prompts,omitempty"`
Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"`
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
Store store.Store `json:"-" yaml:"-"`
RAG *rag.RAG `json:"-" yaml:"-"`
Vision *vision.Vision `json:"-" yaml:"-"`
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
}
// VisionSetting the vision setting
type VisionSetting struct {
Storage driver.StorageConfig `json:"storage" yaml:"storage"`
Model driver.ModelConfig `json:"model" yaml:"model"`
}
// Mention list

208
neo/vision/driver/local/storage.go Normal file
View file

@ -0,0 +1,208 @@
package local
import (
"bytes"
"context"
"crypto/sha256"
"fmt"
"image"
"image/jpeg"
"image/png"
"io"
"path/filepath"
"strings"
"time"
"github.com/yaoapp/gou/fs"
)
// MaxImageSize maximum image size (1920x1080)
const MaxImageSize = 1920
// Storage the local storage driver
type Storage struct {
Path string `json:"path" yaml:"path"`
Compression bool `json:"compression" yaml:"compression"`
BaseURL string `json:"base_url" yaml:"base_url"`
PreviewURL func(fileID string) string `json:"-" yaml:"-"`
}
// New create a new local storage
func New(options map[string]interface{}) (*Storage, error) {
storage := &Storage{
Compression: true,
}
if path, ok := options["path"].(string); ok {
storage.Path = path
}
if compression, ok := options["compression"].(bool); ok {
storage.Compression = compression
}
if baseURL, ok := options["base_url"].(string); ok {
storage.BaseURL = baseURL
}
if previewURL, ok := options["preview_url"].(func(string) string); ok {
storage.PreviewURL = previewURL
}
if storage.Path == "" {
return nil, fmt.Errorf("path is required")
}
return storage, nil
}
// Upload upload file to local storage
func (storage *Storage) Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (string, error) {
data, err := fs.Get("data")
if err != nil {
return "", err
}
ext := filepath.Ext(filename)
id := storage.makeID(filename, ext)
path := filepath.Join(storage.Path, id)
// Create directory if not exists
dir := filepath.Dir(path)
if err := data.MkdirAll(dir, 0755); err != nil {
return "", err
}
// Check if compression is enabled and if it's an image
if storage.Compression && isImage(contentType) {
// Read the entire image into memory
content, err := io.ReadAll(reader)
if err != nil {
return "", fmt.Errorf("failed to read image: %w", err)
}
// Compress image
compressed, err := compressImage(content, contentType)
if err != nil {
return "", fmt.Errorf("failed to compress image: %w", err)
}
// Write compressed image
_, err = data.Write(path, bytes.NewReader(compressed), 0644)
if err != nil {
return "", err
}
} else {
// Write file without compression
_, err = data.Write(path, reader, 0644)
if err != nil {
return "", err
}
}
return id, nil
}
// Download download file from local storage
func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadCloser, string, error) {
data, err := fs.Get("data")
if err != nil {
return nil, "", err
}
path := filepath.Join(storage.Path, fileID)
reader, err := data.ReadCloser(path)
if err != nil {
return nil, "", err
}
contentType := "application/octet-stream"
if v, err := data.MimeType(path); err == nil {
contentType = v
}
return reader, contentType, nil
}
// URL get file url
func (storage *Storage) URL(ctx context.Context, fileID string) string {
if storage.PreviewURL != nil {
return storage.PreviewURL(fileID)
}
if storage.BaseURL != "" {
return fmt.Sprintf("%s/%s", strings.TrimRight(storage.BaseURL, "/"), fileID)
}
return fmt.Sprintf("%s/%s", storage.Path, fileID)
}
func (storage *Storage) makeID(filename string, ext string) string {
date := time.Now().Format("20060102")
hash := fmt.Sprintf("%x", sha256.Sum256([]byte(filename)))[:8]
name := strings.TrimSuffix(filepath.Base(filename), ext)
return fmt.Sprintf("%s/%s-%s%s", date, name, hash, ext)
}
// isImage checks if the content type is an image
func isImage(contentType string) bool {
return strings.HasPrefix(contentType, "image/")
}
// compressImage compresses the image while maintaining aspect ratio
func compressImage(data []byte, contentType string) ([]byte, error) {
// Decode image
img, _, err := image.Decode(bytes.NewReader(data))
if err != nil {
return nil, fmt.Errorf("failed to decode image: %w", err)
}
// Calculate new dimensions
bounds := img.Bounds()
width := bounds.Dx()
height := bounds.Dy()
var newWidth, newHeight int
if width > height {
if width > MaxImageSize {
newWidth = MaxImageSize
newHeight = int(float64(height) * (float64(MaxImageSize) / float64(width)))
} else {
return data, nil // No need to resize
}
} else {
if height > MaxImageSize {
newHeight = MaxImageSize
newWidth = int(float64(width) * (float64(MaxImageSize) / float64(height)))
} else {
return data, nil // No need to resize
}
}
// Create new image with new dimensions
newImg := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
// Scale the image using bilinear interpolation
for y := 0; y < newHeight; y++ {
for x := 0; x < newWidth; x++ {
srcX := float64(x) * float64(width) / float64(newWidth)
srcY := float64(y) * float64(height) / float64(newHeight)
newImg.Set(x, y, img.At(int(srcX), int(srcY)))
}
}
// Encode image
var buf bytes.Buffer
switch contentType {
case "image/jpeg":
err = jpeg.Encode(&buf, newImg, &jpeg.Options{Quality: 85})
case "image/png":
err = png.Encode(&buf, newImg)
default:
return data, nil // Unsupported format, return original
}
if err != nil {
return nil, fmt.Errorf("failed to encode image: %w", err)
}
return buf.Bytes(), nil
}

150
neo/vision/driver/local/storage_test.go Normal file
View file

@ -0,0 +1,150 @@
package local
import (
"bytes"
"context"
"image"
"image/png"
"io"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
func TestLocalStorage(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("Create Storage", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
assert.NotNil(t, storage)
assert.Equal(t, "/__vision_test", storage.Path)
assert.True(t, storage.Compression)
})
t.Run("Upload and Download", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
content := []byte("test content")
reader := bytes.NewReader(content)
fileID, err := storage.Upload(context.Background(), "test.txt", reader, "text/plain")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Contains(t, contentType, "text/plain")
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
assert.Equal(t, content, downloaded)
})
t.Run("Upload and Download Image with Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions
bounds := downloadedImg.Bounds()
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
t.Run("Upload Image without Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": false,
})
assert.NoError(t, err)
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions are unchanged
bounds := downloadedImg.Bounds()
assert.Equal(t, 2000, bounds.Dx())
assert.Equal(t, 2000, bounds.Dy())
})
t.Run("URL Generation", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
fileID := "20240101/test-12345678.txt"
url := storage.URL(context.Background(), fileID)
assert.Equal(t, "/__vision_test/20240101/test-12345678.txt", url)
})
t.Run("Download Non-existent File", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"path": "/__vision_test",
"compression": true,
})
assert.NoError(t, err)
_, _, err = storage.Download(context.Background(), "non-existent.txt")
assert.Error(t, err)
})
}

184
neo/vision/driver/openai/model.go Normal file
View file

@ -0,0 +1,184 @@
package openai
import (
"bytes"
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"github.com/yaoapp/gou/fs"
)
// Model the OpenAI vision model
type Model struct {
APIKey string `json:"api_key" yaml:"api_key"`
Model string `json:"model" yaml:"model"`
Compression bool `json:"compression" yaml:"compression"`
Prompt string `json:"prompt" yaml:"prompt"`
}
// New create a new OpenAI vision model
func New(options map[string]interface{}) (*Model, error) {
model := &Model{
Model: "gpt-4-vision-preview",
Compression: true,
}
if apiKey, ok := options["api_key"].(string); ok {
model.APIKey = apiKey
}
if modelName, ok := options["model"].(string); ok {
model.Model = modelName
}
if compression, ok := options["compression"].(bool); ok {
model.Compression = compression
}
if prompt, ok := options["prompt"].(string); ok {
model.Prompt = prompt
}
if model.APIKey == "" {
return nil, fmt.Errorf("api_key is required")
}
return model, nil
}
// Analyze analyze image using OpenAI vision model
func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error) {
if model.APIKey == "" {
return nil, fmt.Errorf("api_key is required")
}
// Check if fileID is a URL or base64 data
var imageURL string
if strings.HasPrefix(fileID, "data:image/") {
// Already a base64 data URL
imageURL = fileID
} else if strings.HasPrefix(fileID, "http://") || strings.HasPrefix(fileID, "https://") {
// Already a URL
imageURL = fileID
} else {
// Try to read the file and convert to base64
data, err := fs.Get("data")
if err != nil {
return nil, fmt.Errorf("failed to get data fs: %w", err)
}
reader, err := data.ReadCloser(fileID)
if err != nil {
return nil, fmt.Errorf("failed to read file: %w", err)
}
defer reader.Close()
content, err := io.ReadAll(reader)
if err != nil {
return nil, fmt.Errorf("failed to read content: %w", err)
}
// Get content type
contentType := "image/png" // default
if v, err := data.MimeType(fileID); err == nil {
contentType = v
}
// Convert to base64
base64Data := base64.StdEncoding.EncodeToString(content)
imageURL = fmt.Sprintf("data:%s;base64,%s", contentType, base64Data)
}
// Prepare the request body
reqBody := map[string]interface{}{
"model": model.Model,
"messages": []map[string]interface{}{
{
"role": "user",
"content": []map[string]interface{}{
{
"type": "text",
"text": prompt,
},
{
"type": "image_url",
"image_url": map[string]interface{}{
"url": imageURL,
},
},
},
},
},
"max_tokens": 1000,
}
jsonBody, err := json.Marshal(reqBody)
if err != nil {
return nil, fmt.Errorf("failed to marshal request body: %w", err)
}
// Create request
req, err := http.NewRequestWithContext(ctx, "POST", "https://api.openai.com/v1/chat/completions", bytes.NewBuffer(jsonBody))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", model.APIKey))
// Send request
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
// Read response
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("OpenAI API error: %s", string(body))
}
// Parse response
var result map[string]interface{}
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
// Extract content
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("invalid response format")
}
message, ok := choices[0].(map[string]interface{})["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid response format")
}
content, ok := message["content"].(string)
if !ok {
return nil, fmt.Errorf("invalid response format")
}
// Try to parse content as JSON
var description map[string]interface{}
if err := json.Unmarshal([]byte(content), &description); err != nil {
// If not JSON, use the content as description
description = map[string]interface{}{
"description": content,
}
}
return description, nil
}

149
neo/vision/driver/openai/model_test.go Normal file
View file

@ -0,0 +1,149 @@
package openai
import (
"bytes"
"context"
"encoding/base64"
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/neo/vision/driver/s3"
"github.com/yaoapp/yao/test"
)
var (
// 1x1 transparent PNG
testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
func TestOpenAIModel(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("Create Model", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
assert.NotNil(t, model)
if model != nil {
assert.Equal(t, os.Getenv("OPENAI_API_KEY"), model.APIKey)
assert.Equal(t, os.Getenv("VISION_MODEL"), model.Model)
assert.True(t, model.Compression)
}
})
t.Run("Create Model with Invalid API Key", func(t *testing.T) {
_, err := New(map[string]interface{}{})
assert.Error(t, err)
assert.Contains(t, err.Error(), "api_key is required")
})
t.Run("Analyze with Base64 Image", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
// Use base64 image data
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with URL", func(t *testing.T) {
if os.Getenv("S3_API") == "" || os.Getenv("S3_ACCESS_KEY") == "" ||
os.Getenv("S3_SECRET_KEY") == "" || os.Getenv("S3_BUCKET") == "" {
t.Skip("S3 environment variables not set")
}
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
// Create S3 client and upload test image
s3Client, err := s3.New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": "5m",
})
assert.NoError(t, err)
// Upload test image
imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
assert.NoError(t, err)
reader := bytes.NewReader(imgData)
fileID, err := s3Client.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
// Get URL from S3
url := s3Client.URL(context.Background(), fileID)
assert.NotEmpty(t, url)
// Use S3 URL for analysis
result, err := model.Analyze(context.Background(), url, "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with File ID", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
// Create test file
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("/__vision_test/test.png", imgData, 0644)
assert.NoError(t, err)
// Analyze using file ID
result, err := model.Analyze(context.Background(), "/__vision_test/test.png", "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with Invalid File ID", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
_, err = model.Analyze(context.Background(), "/non-existent.png", "Describe this image in detail")
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to read file")
})
t.Run("Analyze with Invalid API Key", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": "invalid-key",
"model": os.Getenv("VISION_MODEL"),
})
assert.NoError(t, err)
_, err = model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Describe this image in detail")
assert.Error(t, err)
assert.Contains(t, err.Error(), "OpenAI API error")
})
}

266
neo/vision/driver/s3/storage.go Normal file
View file

@ -0,0 +1,266 @@
package s3
import (
"bytes"
"context"
"fmt"
"image"
"image/jpeg"
"image/png"
"io"
"path/filepath"
"strings"
"time"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
)
// DefaultExpiration default expiration time for presigned URLs (5 minutes)
const DefaultExpiration = 5 * time.Minute
// MaxImageSize maximum image size (1920x1080)
const MaxImageSize = 1920
// Storage the S3 storage driver
type Storage struct {
Endpoint string `json:"endpoint" yaml:"endpoint"`
Region string `json:"region" yaml:"region"`
Key string `json:"key" yaml:"key"`
Secret string `json:"secret" yaml:"secret"`
Bucket string `json:"bucket" yaml:"bucket"`
Expiration time.Duration `json:"expiration" yaml:"expiration"`
client *s3.Client
prefix string
compression bool
}
// New create a new S3 storage
func New(options map[string]interface{}) (*Storage, error) {
storage := &Storage{
Region: "auto",
Expiration: DefaultExpiration,
compression: true,
}
if endpoint, ok := options["endpoint"].(string); ok {
storage.Endpoint = endpoint
}
if region, ok := options["region"].(string); ok {
storage.Region = region
}
if key, ok := options["key"].(string); ok {
storage.Key = key
}
if secret, ok := options["secret"].(string); ok {
storage.Secret = secret
}
if bucket, ok := options["bucket"].(string); ok {
storage.Bucket = bucket
}
if prefix, ok := options["prefix"].(string); ok {
storage.prefix = prefix
}
if exp, ok := options["expiration"].(time.Duration); ok {
storage.Expiration = exp
}
if compression, ok := options["compression"].(bool); ok {
storage.compression = compression
}
// Validate required fields
if storage.Key == "" || storage.Secret == "" {
return nil, fmt.Errorf("key and secret are required")
}
if storage.Bucket == "" {
return nil, fmt.Errorf("bucket is required")
}
// Create S3 client
opts := s3.Options{
Region: storage.Region,
Credentials: credentials.NewStaticCredentialsProvider(storage.Key, storage.Secret, ""),
UsePathStyle: true,
}
if storage.Endpoint != "" {
// Remove bucket name from endpoint if present
endpoint := storage.Endpoint
if strings.Contains(endpoint, "/"+storage.Bucket) {
endpoint = strings.TrimSuffix(endpoint, "/"+storage.Bucket)
}
opts.BaseEndpoint = aws.String(endpoint)
}
storage.client = s3.New(opts)
return storage, nil
}
// Upload upload file to S3
func (storage *Storage) Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (string, error) {
if storage.client == nil {
return "", fmt.Errorf("s3 client not initialized")
}
// Generate file ID
fileID := storage.makeID(filename, filepath.Ext(filename))
key := filepath.Join(storage.prefix, fileID)
// Check if compression is enabled and if it's an image
var body io.Reader
if storage.compression && isImage(contentType) {
// Read the entire image into memory
content, err := io.ReadAll(reader)
if err != nil {
return "", fmt.Errorf("failed to read image: %w", err)
}
// Compress image
compressed, err := compressImage(content, contentType)
if err != nil {
return "", fmt.Errorf("failed to compress image: %w", err)
}
body = bytes.NewReader(compressed)
} else {
body = reader
}
// Upload file
_, err := storage.client.PutObject(ctx, &s3.PutObjectInput{
Bucket: aws.String(storage.Bucket),
Key: aws.String(key),
Body: body,
ContentType: aws.String(contentType),
})
if err != nil {
return "", fmt.Errorf("failed to upload file: %w", err)
}
return fileID, nil
}
// Download download file from S3
func (storage *Storage) Download(ctx context.Context, fileID string) (io.ReadCloser, string, error) {
if storage.client == nil {
return nil, "", fmt.Errorf("s3 client not initialized")
}
key := filepath.Join(storage.prefix, fileID)
// Get object
result, err := storage.client.GetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(storage.Bucket),
Key: aws.String(key),
})
if err != nil {
return nil, "", fmt.Errorf("failed to download file: %w", err)
}
contentType := "application/octet-stream"
if result.ContentType != nil {
contentType = *result.ContentType
}
return result.Body, contentType, nil
}
// URL get file url with expiration
func (storage *Storage) URL(ctx context.Context, fileID string) string {
if storage.client == nil {
return ""
}
key := filepath.Join(storage.prefix, fileID)
presignClient := s3.NewPresignClient(storage.client)
request, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
Bucket: aws.String(storage.Bucket),
Key: aws.String(key),
}, s3.WithPresignExpires(storage.Expiration))
if err != nil {
return ""
}
return request.URL
}
func (storage *Storage) makeID(filename string, ext string) string {
date := time.Now().Format("20060102")
name := strings.TrimSuffix(filepath.Base(filename), ext)
return fmt.Sprintf("%s/%s-%d%s", date, name, time.Now().UnixNano(), ext)
}
// isImage checks if the content type is an image
func isImage(contentType string) bool {
return strings.HasPrefix(contentType, "image/")
}
// compressImage compresses the image while maintaining aspect ratio
func compressImage(data []byte, contentType string) ([]byte, error) {
// Decode image
img, _, err := image.Decode(bytes.NewReader(data))
if err != nil {
return nil, fmt.Errorf("failed to decode image: %w", err)
}
// Calculate new dimensions
bounds := img.Bounds()
width := bounds.Dx()
height := bounds.Dy()
var newWidth, newHeight int
if width > height {
if width > MaxImageSize {
newWidth = MaxImageSize
newHeight = int(float64(height) * (float64(MaxImageSize) / float64(width)))
} else {
return data, nil // No need to resize
}
} else {
if height > MaxImageSize {
newHeight = MaxImageSize
newWidth = int(float64(width) * (float64(MaxImageSize) / float64(height)))
} else {
return data, nil // No need to resize
}
}
// Create new image with new dimensions
newImg := image.NewRGBA(image.Rect(0, 0, newWidth, newHeight))
// Scale the image using bilinear interpolation
for y := 0; y < newHeight; y++ {
for x := 0; x < newWidth; x++ {
srcX := float64(x) * float64(width) / float64(newWidth)
srcY := float64(y) * float64(height) / float64(newHeight)
newImg.Set(x, y, img.At(int(srcX), int(srcY)))
}
}
// Encode image
var buf bytes.Buffer
switch contentType {
case "image/jpeg":
err = jpeg.Encode(&buf, newImg, &jpeg.Options{Quality: 85})
case "image/png":
err = png.Encode(&buf, newImg)
default:
return data, nil // Unsupported format, return original
}
if err != nil {
return nil, fmt.Errorf("failed to encode image: %w", err)
}
return buf.Bytes(), nil
}

204
neo/vision/driver/s3/storage_test.go Normal file
View file

@ -0,0 +1,204 @@
package s3
import (
"bytes"
"context"
"image"
"image/png"
"io"
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
func TestS3Storage(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("Create Storage", func(t *testing.T) {
options := map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 10 * time.Minute,
"compression": true,
}
storage, err := New(options)
if err != nil {
t.Logf("Error creating storage: %v", err)
}
assert.NoError(t, err)
assert.NotNil(t, storage)
if storage != nil {
assert.Equal(t, os.Getenv("S3_API"), storage.Endpoint)
assert.Equal(t, "auto", storage.Region)
assert.Equal(t, os.Getenv("S3_ACCESS_KEY"), storage.Key)
assert.Equal(t, os.Getenv("S3_SECRET_KEY"), storage.Secret)
assert.Equal(t, os.Getenv("S3_BUCKET"), storage.Bucket)
assert.Equal(t, "vision-test", storage.prefix)
assert.Equal(t, 10*time.Minute, storage.Expiration)
assert.True(t, storage.compression)
}
})
t.Run("Upload and Download Image with Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": true,
})
if err != nil {
t.Skip("S3 configuration not available")
}
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions
bounds := downloadedImg.Bounds()
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
t.Run("Upload Image without Compression", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": false,
})
if err != nil {
t.Skip("S3 configuration not available")
}
// Create a test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
fileID, err := storage.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Download and verify size
reader2, contentType, err := storage.Download(context.Background(), fileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions are unchanged
bounds := downloadedImg.Bounds()
assert.Equal(t, 2000, bounds.Dx())
assert.Equal(t, 2000, bounds.Dy())
})
t.Run("Upload and Download Text File", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": true,
})
if err != nil {
t.Skip("S3 configuration not available")
}
content := []byte("test content")
reader := bytes.NewReader(content)
fileID, err := storage.Upload(context.Background(), "test.txt", reader, "text/plain")
assert.NoError(t, err)
assert.NotEmpty(t, fileID)
// Get presigned URL
url := storage.URL(context.Background(), fileID)
assert.NotEmpty(t, url)
assert.Contains(t, url, "X-Amz-Signature")
assert.Contains(t, url, "X-Amz-Expires")
// Download
reader2, contentType, err := storage.Download(context.Background(), fileID)
if err != nil {
t.Logf("Download error: %v", err)
t.FailNow()
}
assert.NoError(t, err)
assert.Contains(t, contentType, "text/plain")
if reader2 != nil {
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
assert.Equal(t, content, downloaded)
reader2.Close()
}
})
t.Run("Download Non-existent File", func(t *testing.T) {
storage, err := New(map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": 5 * time.Minute,
"compression": true,
})
if err != nil {
t.Skip("S3 configuration not available")
}
_, _, err = storage.Download(context.Background(), "non-existent.txt")
assert.Error(t, err)
})
}

43
neo/vision/driver/types.go Normal file
View file

@ -0,0 +1,43 @@
package driver
import (
"context"
"io"
)
// Config the vision configuration
type Config struct {
Storage StorageConfig `json:"storage" yaml:"storage"`
Model ModelConfig `json:"model" yaml:"model"`
}
// StorageConfig the storage configuration
type StorageConfig struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// ModelConfig the model configuration
type ModelConfig struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// Storage the storage interface
type Storage interface {
Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (string, error)
Download(ctx context.Context, fileID string) (io.ReadCloser, string, error)
URL(ctx context.Context, fileID string) string
}
// Model the vision model interface
type Model interface {
Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error)
}
// Response the vision response
type Response struct {
FileID string `json:"file_id" yaml:"file_id"`
URL string `json:"url" yaml:"url"`
Description map[string]interface{} `json:"description" yaml:"description"`
}

110
neo/vision/vision.go Normal file
View file

@ -0,0 +1,110 @@
package vision
import (
"context"
"fmt"
"io"
"strings"
"time"
"github.com/yaoapp/yao/neo/vision/driver"
"github.com/yaoapp/yao/neo/vision/driver/local"
"github.com/yaoapp/yao/neo/vision/driver/openai"
"github.com/yaoapp/yao/neo/vision/driver/s3"
)
// Vision the vision service
type Vision struct {
storage driver.Storage
model driver.Model
}
// New create a new vision service
func New(cfg *driver.Config) (*Vision, error) {
// Create storage driver
var storage driver.Storage
var err error
switch cfg.Storage.Driver {
case "local":
storage, err = local.New(cfg.Storage.Options)
case "s3":
// Convert expiration string to duration if present
if exp, ok := cfg.Storage.Options["expiration"].(string); ok {
if duration, err := time.ParseDuration(exp); err == nil {
cfg.Storage.Options["expiration"] = duration
}
}
storage, err = s3.New(cfg.Storage.Options)
default:
return nil, fmt.Errorf("storage driver %s not supported", cfg.Storage.Driver)
}
if err != nil {
return nil, fmt.Errorf("create storage driver error: %s", err.Error())
}
// Create model driver
var model driver.Model
switch cfg.Model.Driver {
case "openai":
model, err = openai.New(cfg.Model.Options)
default:
return nil, fmt.Errorf("model driver %s not supported", cfg.Model.Driver)
}
if err != nil {
return nil, fmt.Errorf("create model driver error: %s", err.Error())
}
return &Vision{
storage: storage,
model: model,
}, nil
}
// Upload upload file
func (v *Vision) Upload(ctx context.Context, filename string, reader io.Reader, contentType string) (*driver.Response, error) {
fileID, err := v.storage.Upload(ctx, filename, reader, contentType)
if err != nil {
return nil, err
}
return &driver.Response{
FileID: fileID,
URL: v.storage.URL(ctx, fileID),
}, nil
}
// Analyze analyze image using vision model
func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*driver.Response, error) {
if v.model == nil {
return nil, fmt.Errorf("model is required")
}
var url string
// If the input is already a base64 data URL or a HTTP(S) URL, use it directly
if strings.HasPrefix(fileID, "data:image/") || strings.HasPrefix(fileID, "http://") || strings.HasPrefix(fileID, "https://") {
url = fileID
} else {
// Otherwise, try to get the URL from storage
url = v.storage.URL(ctx, fileID)
if url == "" {
return nil, fmt.Errorf("failed to get URL for file %s", fileID)
}
}
result, err := v.model.Analyze(ctx, url, prompt)
if err != nil {
return nil, err
}
return &driver.Response{
FileID: fileID,
URL: url,
Description: result,
}, nil
}
// Download download file
func (v *Vision) Download(ctx context.Context, fileID string) (io.ReadCloser, string, error) {
return v.storage.Download(ctx, fileID)
}

412
neo/vision/vision_test.go Normal file
View file

@ -0,0 +1,412 @@
package vision
import (
"bytes"
"context"
"encoding/base64"
"fmt"
"image"
"image/png"
"io"
"net/http"
"net/http/httptest"
"os"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/neo/vision/driver"
"github.com/yaoapp/yao/neo/vision/driver/local"
"github.com/yaoapp/yao/test"
)
var (
// 1x1 transparent PNG
testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
)
// MaxImageSize maximum image size (1920x1080)
const MaxImageSize = local.MaxImageSize
func TestVision(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Setup test server for image hosting
imgServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Log request for debugging
t.Logf("Received request for: %s", r.URL.Path)
// Always return the test image
imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
w.Header().Set("Content-Type", "image/png")
w.Write(imgData)
}))
defer imgServer.Close()
t.Logf("Test server running at: %s", imgServer.URL)
t.Run("Create Vision Service", func(t *testing.T) {
vision, err := createTestVision(imgServer.URL)
assert.NoError(t, err)
assert.NotNil(t, vision)
})
t.Run("Upload and Download with Local Storage", func(t *testing.T) {
vision, err := createTestVision(imgServer.URL)
assert.NoError(t, err)
// Test with text file
content := []byte("test content")
reader := bytes.NewReader(content)
resp, err := vision.Upload(context.Background(), "test.txt", reader, "text/plain")
assert.NoError(t, err)
assert.NotEmpty(t, resp.FileID)
assert.NotEmpty(t, resp.URL)
// Download
reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
assert.NoError(t, err)
assert.Contains(t, contentType, "text/plain")
if reader2 != nil {
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
assert.Equal(t, content, downloaded)
reader2.Close()
}
})
t.Run("Upload and Download with S3 Storage", func(t *testing.T) {
vision, err := createTestVisionWithS3()
if err != nil {
t.Skip("S3 configuration not available")
}
// Test with text file
content := []byte("test content")
reader := bytes.NewReader(content)
resp, err := vision.Upload(context.Background(), "test.txt", reader, "text/plain")
assert.NoError(t, err)
assert.NotEmpty(t, resp.FileID)
assert.NotEmpty(t, resp.URL)
// Download
reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
assert.NoError(t, err)
assert.Contains(t, contentType, "text/plain")
if reader2 != nil {
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
assert.Equal(t, content, downloaded)
reader2.Close()
}
})
t.Run("Analyze Image with Base64", func(t *testing.T) {
// Create vision service
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "local",
Options: map[string]interface{}{
"path": "/__vision_test",
"compression": true,
},
},
Model: driver.ModelConfig{
Driver: "openai",
Options: map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
},
},
}
vision, err := New(cfg)
assert.NoError(t, err)
// Use base64 data directly
result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Description)
})
t.Run("Analyze Image with File", func(t *testing.T) {
// Create vision service
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "local",
Options: map[string]interface{}{
"path": "/__vision_test",
"compression": true,
},
},
Model: driver.ModelConfig{
Driver: "openai",
Options: map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
},
},
}
vision, err := New(cfg)
assert.NoError(t, err)
// Create test file
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)
// Analyze using file path
result, err := vision.Analyze(context.Background(), "/test.png", "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Description)
})
t.Run("Analyze Image with S3 URL", func(t *testing.T) {
if os.Getenv("S3_API") == "" || os.Getenv("S3_ACCESS_KEY") == "" ||
os.Getenv("S3_SECRET_KEY") == "" || os.Getenv("S3_BUCKET") == "" {
t.Skip("S3 environment variables not set")
}
// Create vision service
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "s3",
Options: map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": "5m",
},
},
Model: driver.ModelConfig{
Driver: "openai",
Options: map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
},
},
}
vision, err := New(cfg)
assert.NoError(t, err)
// Upload test image
imgData, err := base64.StdEncoding.DecodeString(testImageBase64)
assert.NoError(t, err)
reader := bytes.NewReader(imgData)
resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, resp.FileID)
assert.NotEmpty(t, resp.URL)
// Analyze using S3 URL
result, err := vision.Analyze(context.Background(), resp.URL, "Describe this image in detail")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Description)
})
t.Run("Invalid Model", func(t *testing.T) {
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "local",
Options: map[string]interface{}{
"path": "/__vision_test",
"compression": true,
},
},
Model: driver.ModelConfig{
Driver: "invalid",
Options: map[string]interface{}{},
},
}
_, err := New(cfg)
assert.Error(t, err)
assert.Contains(t, err.Error(), "model driver invalid not supported")
})
t.Run("Invalid Storage", func(t *testing.T) {
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "invalid",
Options: map[string]interface{}{},
},
Model: driver.ModelConfig{
Driver: "openai",
Options: map[string]interface{}{
"api_key": "test",
},
},
}
_, err := New(cfg)
assert.Error(t, err)
assert.Contains(t, err.Error(), "storage driver invalid not supported")
})
t.Run("Upload and Download Image with Local Storage", func(t *testing.T) {
vision, err := createTestVision(imgServer.URL)
assert.NoError(t, err)
// Create test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, resp.FileID)
assert.NotEmpty(t, resp.URL)
// Download and verify size
reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions
bounds := downloadedImg.Bounds()
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
t.Run("Upload and Download Image with S3 Storage", func(t *testing.T) {
vision, err := createTestVisionWithS3()
if err != nil {
t.Skip("S3 configuration not available")
}
// Create test image (2000x2000 pixels)
img := image.NewRGBA(image.Rect(0, 0, 2000, 2000))
var buf bytes.Buffer
err = png.Encode(&buf, img)
assert.NoError(t, err)
// Upload
reader := bytes.NewReader(buf.Bytes())
resp, err := vision.Upload(context.Background(), "test.png", reader, "image/png")
assert.NoError(t, err)
assert.NotEmpty(t, resp.FileID)
assert.NotEmpty(t, resp.URL)
// Download and verify size
reader2, contentType, err := vision.Download(context.Background(), resp.FileID)
assert.NoError(t, err)
assert.Equal(t, "image/png", contentType)
downloaded, err := io.ReadAll(reader2)
assert.NoError(t, err)
// Decode the downloaded image
downloadedImg, _, err := image.Decode(bytes.NewReader(downloaded))
assert.NoError(t, err)
// Verify dimensions
bounds := downloadedImg.Bounds()
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
}
func createTestVision(baseURL string) (*Vision, error) {
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "local",
Options: map[string]interface{}{
"path": "/__vision_test",
"compression": true,
"base_url": baseURL,
},
},
Model: driver.ModelConfig{
Driver: "openai",
Options: map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": `# Objective
You are a vision assistant, you can help the user to understand the image and describe it.
## Task Execution Steps
1. Understand the image/video and describe it.
2. Describe the image/video in detail.
## Result Format
{
"description": "The description of the image/video",
"content": "The content of the image/video"
}`,
},
},
}
return New(cfg)
}
func createTestVisionWithS3() (*Vision, error) {
// Check required S3 environment variables
if os.Getenv("S3_API") == "" || os.Getenv("S3_ACCESS_KEY") == "" ||
os.Getenv("S3_SECRET_KEY") == "" || os.Getenv("S3_BUCKET") == "" {
return nil, fmt.Errorf("S3 environment variables not set")
}
cfg := &driver.Config{
Storage: driver.StorageConfig{
Driver: "s3",
Options: map[string]interface{}{
"endpoint": os.Getenv("S3_API"),
"region": "auto",
"key": os.Getenv("S3_ACCESS_KEY"),
"secret": os.Getenv("S3_SECRET_KEY"),
"bucket": os.Getenv("S3_BUCKET"),
"prefix": "vision-test",
"expiration": "5m",
},
},
Model: driver.ModelConfig{
Driver: "openai",
Options: map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": `# Objective
You are a vision assistant, you can help the user to understand the image and describe it.
## Task Execution Steps
1. Understand the image/video and describe it.
2. Describe the image/video in detail.
## Result Format
{
"description": "The description of the image/video",
"content": "The content of the image/video"
}`,
},
},
}
return New(cfg)
}