From 93a6d7a6a09b5b9877825ffb6818e42afa6d784e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 3 Jan 2025 13:33:58 +0800 Subject: [PATCH 1/5] Update Go module dependencies to include AWS SDK v2 - Added several AWS SDK v2 modules to go.mod and go.sum, enhancing cloud service integration capabilities: - github.com/aws/aws-sdk-go-v2 and its related services (S3, STS, SSO, etc.) for improved AWS functionality. - Updated indirect dependencies to ensure compatibility and performance improvements. - This update supports future development involving AWS services, streamlining the integration process. --- go.mod | 18 ++ go.sum | 36 +++ neo/vision/driver/local/storage.go | 115 ++++++++ neo/vision/driver/local/storage_test.go | 74 ++++++ neo/vision/driver/openai/model.go | 184 +++++++++++++ neo/vision/driver/openai/model_test.go | 149 +++++++++++ neo/vision/driver/s3/storage.go | 168 ++++++++++++ neo/vision/driver/s3/storage_test.go | 143 ++++++++++ neo/vision/driver/types.go | 43 +++ neo/vision/prompt.md | 8 + neo/vision/vision.go | 110 ++++++++ neo/vision/vision_test.go | 334 ++++++++++++++++++++++++ 12 files changed, 1382 insertions(+) create mode 100644 neo/vision/driver/local/storage.go create mode 100644 neo/vision/driver/local/storage_test.go create mode 100644 neo/vision/driver/openai/model.go create mode 100644 neo/vision/driver/openai/model_test.go create mode 100644 neo/vision/driver/s3/storage.go create mode 100644 neo/vision/driver/s3/storage_test.go create mode 100644 neo/vision/driver/types.go create mode 100644 neo/vision/prompt.md create mode 100644 neo/vision/vision.go create mode 100644 neo/vision/vision_test.go diff --git a/go.mod b/go.mod index e99535e3..f3c33bf8 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,10 @@ 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/config v1.28.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 +42,20 @@ 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/feature/ec2/imds v1.16.22 // 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/ini v1.8.1 // 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/aws-sdk-go-v2/service/sso v1.24.8 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.33.3 // 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 diff --git a/go.sum b/go.sum index a9b45ed1..a26a9c8b 100644 --- a/go.sum +++ b/go.sum @@ -6,6 +6,42 @@ 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/config v1.28.7 h1:GduUnoTXlhkgnxTD93g1nv4tVPILbdNQOzav+Wpg7AE= +github.com/aws/aws-sdk-go-v2/config v1.28.7/go.mod h1:vZGX6GVkIE8uECSUHB6MWAUsd4ZcG2Yq/dMa4refR3M= +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/feature/ec2/imds v1.16.22 h1:kqOrpojG71DxJm/KDPO+Z/y1phm1JlC8/iT+5XRmAn8= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.22/go.mod h1:NtSFajXVVL8TA2QNngagVZmUtXciyrHOt7xgz4faS/M= +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/ini v1.8.1 h1:VaRN3TlFdd6KxX1x3ILT5ynH6HvKgqdiXoTxAF4HQcQ= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc= +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/aws-sdk-go-v2/service/sso v1.24.8 h1:CvuUmnXI7ebaUAhbJcDy9YQx8wHR69eZ9I7q5hszt/g= +github.com/aws/aws-sdk-go-v2/service/sso v1.24.8/go.mod h1:XDeGv1opzwm8ubxddF0cgqkZWsyOtw4lr6dxwmb6YQg= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7 h1:F2rBfNAL5UyswqoeWv9zs74N/NanhK16ydHW1pahX6E= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7/go.mod h1:JfyQ0g2JG8+Krq0EuZNnRwX0mU0HrwY/tG6JNfcqh4k= +github.com/aws/aws-sdk-go-v2/service/sts v1.33.3 h1:Xgv/hyNgvLda/M9l9qxXc4UFSgppnRczLxlMs5Ae/QY= +github.com/aws/aws-sdk-go-v2/service/sts v1.33.3/go.mod h1:5Gn+d+VaaRgsjewpMvGazt0WfcFO+Md4wLOuBfGR9Bc= +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= diff --git a/neo/vision/driver/local/storage.go b/neo/vision/driver/local/storage.go new file mode 100644 index 00000000..ed873129 --- /dev/null +++ b/neo/vision/driver/local/storage.go @@ -0,0 +1,115 @@ +package local + +import ( + "context" + "crypto/sha256" + "fmt" + "io" + "path/filepath" + "strings" + "time" + + "github.com/yaoapp/gou/fs" +) + +// 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 + } + + // Write file + _, 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) +} diff --git a/neo/vision/driver/local/storage_test.go b/neo/vision/driver/local/storage_test.go new file mode 100644 index 00000000..15baa1a2 --- /dev/null +++ b/neo/vision/driver/local/storage_test.go @@ -0,0 +1,74 @@ +package local + +import ( + "bytes" + "context" + "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("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) + }) +} diff --git a/neo/vision/driver/openai/model.go b/neo/vision/driver/openai/model.go new file mode 100644 index 00000000..a766c380 --- /dev/null +++ b/neo/vision/driver/openai/model.go @@ -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 +} diff --git a/neo/vision/driver/openai/model_test.go b/neo/vision/driver/openai/model_test.go new file mode 100644 index 00000000..e93d516e --- /dev/null +++ b/neo/vision/driver/openai/model_test.go @@ -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") + }) +} diff --git a/neo/vision/driver/s3/storage.go b/neo/vision/driver/s3/storage.go new file mode 100644 index 00000000..addc3b10 --- /dev/null +++ b/neo/vision/driver/s3/storage.go @@ -0,0 +1,168 @@ +package s3 + +import ( + "context" + "fmt" + "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 + +// 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 +} + +// New create a new S3 storage +func New(options map[string]interface{}) (*Storage, error) { + storage := &Storage{ + Region: "auto", + Expiration: DefaultExpiration, + } + + 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 + } + + // 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) + + // Upload file + _, err := storage.client.PutObject(ctx, &s3.PutObjectInput{ + Bucket: aws.String(storage.Bucket), + Key: aws.String(key), + Body: reader, + 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) +} diff --git a/neo/vision/driver/s3/storage_test.go b/neo/vision/driver/s3/storage_test.go new file mode 100644 index 00000000..0f3d3847 --- /dev/null +++ b/neo/vision/driver/s3/storage_test.go @@ -0,0 +1,143 @@ +package s3 + +import ( + "bytes" + "context" + "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, + } + + 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) + } + }) + + t.Run("Upload and Download", 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, + }) + 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) + + // 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") + + // Test with different expiration + storage.Expiration = 1 * time.Hour + url2 := storage.URL(context.Background(), fileID) + assert.NotEmpty(t, url2) + assert.Contains(t, url2, "X-Amz-Signature") + assert.Contains(t, url2, "X-Amz-Expires=3600") + + // 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("Upload with Custom Expiration", 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, + }) + 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) + + // Get URL with default expiration (5 minutes) + url := storage.URL(context.Background(), fileID) + assert.NotEmpty(t, url) + assert.Contains(t, url, "X-Amz-Signature") + assert.Contains(t, url, "X-Amz-Expires=300") // 5 minutes = 300 seconds + + // Change expiration and get new URL + storage.Expiration = 2 * time.Hour + url2 := storage.URL(context.Background(), fileID) + assert.NotEmpty(t, url2) + assert.Contains(t, url2, "X-Amz-Signature") + assert.Contains(t, url2, "X-Amz-Expires=7200") // 2 hours = 7200 seconds + }) + + 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, + }) + assert.NoError(t, err) + + _, _, err = storage.Download(context.Background(), "non-existent.txt") + assert.Error(t, err) + }) +} diff --git a/neo/vision/driver/types.go b/neo/vision/driver/types.go new file mode 100644 index 00000000..ccc64ee6 --- /dev/null +++ b/neo/vision/driver/types.go @@ -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"` +} diff --git a/neo/vision/prompt.md b/neo/vision/prompt.md new file mode 100644 index 00000000..2d6e4f34 --- /dev/null +++ b/neo/vision/prompt.md @@ -0,0 +1,8 @@ +根据这个数据结构，和说明实现一下对应的逻辑。 + +1. driver 单独一个目录. 每个 driver 一个目录，model 和 storage 在一级即可。 放在@vision 下 +2. model driver: 支持 openai +3. storage driver 支持 local 和 s3 . local 使用我框架的 fs 实现，参考@file.go +4. 在 程序启动时候，设置视觉配置。(作为可选配置） @types.go +5. 统一的创建和调用入口，调用时需传入 chat model (用来判断是否支持视觉) 和图片路径，（使用 fs 读取）。 @vision.go +6. 用一个回调函数，外部传入用来格式化返回的数据。 diff --git a/neo/vision/vision.go b/neo/vision/vision.go new file mode 100644 index 00000000..79c5f372 --- /dev/null +++ b/neo/vision/vision.go @@ -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) +} diff --git a/neo/vision/vision_test.go b/neo/vision/vision_test.go new file mode 100644 index 00000000..73379487 --- /dev/null +++ b/neo/vision/vision_test.go @@ -0,0 +1,334 @@ +package vision + +import ( + "bytes" + "context" + "encoding/base64" + "fmt" + "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/test" +) + +var ( + // 1x1 transparent PNG + testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" +) + +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") + }) +} + +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) +} From 1ea18526774887d329143050b1a9c60c77eb14dd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 3 Jan 2025 13:52:45 +0800 Subject: [PATCH 2/5] Remove prompt documentation and enhance image upload/download functionality in vision module - Deleted the obsolete prompt.md file, streamlining the codebase. - Added tests for image upload and download functionality using local and S3 storage drivers, ensuring robust handling of image files. - Implemented image compression during upload for both local and S3 storage, maintaining aspect ratio and limiting dimensions to a maximum size of 1920x1080. - Enhanced the local and S3 storage drivers to support image compression, improving storage efficiency. - Updated tests to verify image dimensions and content type after upload and download operations, ensuring data integrity. --- neo/vision/driver/local/storage.go | 101 ++++++++++++- neo/vision/driver/local/storage_test.go | 76 ++++++++++ neo/vision/driver/s3/storage.go | 120 +++++++++++++-- neo/vision/driver/s3/storage_test.go | 185 ++++++++++++++++-------- neo/vision/prompt.md | 8 - neo/vision/vision_test.go | 100 +++++++++++-- 6 files changed, 494 insertions(+), 96 deletions(-) delete mode 100644 neo/vision/prompt.md diff --git a/neo/vision/driver/local/storage.go b/neo/vision/driver/local/storage.go index ed873129..81de81a3 100644 --- a/neo/vision/driver/local/storage.go +++ b/neo/vision/driver/local/storage.go @@ -1,9 +1,13 @@ package local import ( + "bytes" "context" "crypto/sha256" "fmt" + "image" + "image/jpeg" + "image/png" "io" "path/filepath" "strings" @@ -12,6 +16,9 @@ import ( "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"` @@ -66,10 +73,31 @@ func (storage *Storage) Upload(ctx context.Context, filename string, reader io.R return "", err } - // Write file - _, err = data.Write(path, reader, 0644) - if 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 @@ -113,3 +141,68 @@ func (storage *Storage) makeID(filename string, ext string) string { 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 +} diff --git a/neo/vision/driver/local/storage_test.go b/neo/vision/driver/local/storage_test.go index 15baa1a2..36a53eee 100644 --- a/neo/vision/driver/local/storage_test.go +++ b/neo/vision/driver/local/storage_test.go @@ -3,6 +3,8 @@ package local import ( "bytes" "context" + "image" + "image/png" "io" "testing" @@ -49,6 +51,80 @@ func TestLocalStorage(t *testing.T) { 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", diff --git a/neo/vision/driver/s3/storage.go b/neo/vision/driver/s3/storage.go index addc3b10..3f55b7be 100644 --- a/neo/vision/driver/s3/storage.go +++ b/neo/vision/driver/s3/storage.go @@ -1,8 +1,12 @@ package s3 import ( + "bytes" "context" "fmt" + "image" + "image/jpeg" + "image/png" "io" "path/filepath" "strings" @@ -16,23 +20,28 @@ import ( // 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 + 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, + Region: "auto", + Expiration: DefaultExpiration, + compression: true, } if endpoint, ok := options["endpoint"].(string); ok { @@ -63,6 +72,10 @@ func New(options map[string]interface{}) (*Storage, error) { 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") @@ -102,11 +115,31 @@ func (storage *Storage) Upload(ctx context.Context, filename string, reader io.R 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: reader, + Body: body, ContentType: aws.String(contentType), }) if err != nil { @@ -166,3 +199,68 @@ func (storage *Storage) makeID(filename string, ext string) string { 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 +} diff --git a/neo/vision/driver/s3/storage_test.go b/neo/vision/driver/s3/storage_test.go index 0f3d3847..ca0737ce 100644 --- a/neo/vision/driver/s3/storage_test.go +++ b/neo/vision/driver/s3/storage_test.go @@ -3,6 +3,8 @@ package s3 import ( "bytes" "context" + "image" + "image/png" "io" "os" "testing" @@ -19,13 +21,14 @@ func TestS3Storage(t *testing.T) { 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, + "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) @@ -42,21 +45,115 @@ func TestS3Storage(t *testing.T) { 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", func(t *testing.T) { + 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, + "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") @@ -69,13 +166,6 @@ func TestS3Storage(t *testing.T) { assert.Contains(t, url, "X-Amz-Signature") assert.Contains(t, url, "X-Amz-Expires") - // Test with different expiration - storage.Expiration = 1 * time.Hour - url2 := storage.URL(context.Background(), fileID) - assert.NotEmpty(t, url2) - assert.Contains(t, url2, "X-Amz-Signature") - assert.Contains(t, url2, "X-Amz-Expires=3600") - // Download reader2, contentType, err := storage.Download(context.Background(), fileID) if err != nil { @@ -93,49 +183,20 @@ func TestS3Storage(t *testing.T) { } }) - t.Run("Upload with Custom Expiration", 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, - }) - 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) - - // Get URL with default expiration (5 minutes) - url := storage.URL(context.Background(), fileID) - assert.NotEmpty(t, url) - assert.Contains(t, url, "X-Amz-Signature") - assert.Contains(t, url, "X-Amz-Expires=300") // 5 minutes = 300 seconds - - // Change expiration and get new URL - storage.Expiration = 2 * time.Hour - url2 := storage.URL(context.Background(), fileID) - assert.NotEmpty(t, url2) - assert.Contains(t, url2, "X-Amz-Signature") - assert.Contains(t, url2, "X-Amz-Expires=7200") // 2 hours = 7200 seconds - }) - 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, + "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, }) - assert.NoError(t, err) + if err != nil { + t.Skip("S3 configuration not available") + } _, _, err = storage.Download(context.Background(), "non-existent.txt") assert.Error(t, err) diff --git a/neo/vision/prompt.md b/neo/vision/prompt.md deleted file mode 100644 index 2d6e4f34..00000000 --- a/neo/vision/prompt.md +++ /dev/null @@ -1,8 +0,0 @@ -根据这个数据结构，和说明实现一下对应的逻辑。 - -1. driver 单独一个目录. 每个 driver 一个目录，model 和 storage 在一级即可。 放在@vision 下 -2. model driver: 支持 openai -3. storage driver 支持 local 和 s3 . local 使用我框架的 fs 实现，参考@file.go -4. 在 程序启动时候，设置视觉配置。(作为可选配置） @types.go -5. 统一的创建和调用入口，调用时需传入 chat model (用来判断是否支持视觉) 和图片路径，（使用 fs 读取）。 @vision.go -6. 用一个回调函数，外部传入用来格式化返回的数据。 diff --git a/neo/vision/vision_test.go b/neo/vision/vision_test.go index 73379487..a0bd42ae 100644 --- a/neo/vision/vision_test.go +++ b/neo/vision/vision_test.go @@ -5,6 +5,8 @@ import ( "context" "encoding/base64" "fmt" + "image" + "image/png" "io" "net/http" "net/http/httptest" @@ -15,6 +17,7 @@ import ( "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" ) @@ -23,6 +26,9 @@ var ( 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() @@ -253,6 +259,78 @@ func TestVision(t *testing.T) { 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) { @@ -271,17 +349,17 @@ func createTestVision(baseURL string) (*Vision, error) { "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" - }`, + 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" + }`, }, }, } From 966b0f601f4db8eb0ffd684d2e27413b2e6152e1 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 3 Jan 2025 15:56:14 +0800 Subject: [PATCH 3/5] Update Go module dependencies and enhance Neo API with Vision integration - Removed obsolete AWS SDK v2 dependencies from go.mod and go.sum to streamline the codebase. - Added Vision-related imports and initialized Vision settings in the Neo API, enhancing AI capabilities. - Introduced initVision method to configure and initialize the Vision instance, improving modularity and functionality. - Updated the DSL struct to include Vision settings, allowing for better configuration management. This commit supports future enhancements in AI functionalities and improves the overall structure of the Neo API. --- go.mod | 6 ------ go.sum | 12 ------------ neo/load.go | 26 ++++++++++++++++++++++++++ neo/types.go | 10 ++++++++++ 4 files changed, 36 insertions(+), 18 deletions(-) diff --git a/go.mod b/go.mod index f3c33bf8..082cb348 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,6 @@ 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/config v1.28.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 @@ -43,18 +42,13 @@ require ( 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/feature/ec2/imds v1.16.22 // 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/ini v1.8.1 // 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/aws-sdk-go-v2/service/sso v1.24.8 // indirect - github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7 // indirect - github.com/aws/aws-sdk-go-v2/service/sts v1.33.3 // 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 diff --git a/go.sum b/go.sum index a26a9c8b..09a2c82a 100644 --- a/go.sum +++ b/go.sum @@ -10,18 +10,12 @@ github.com/aws/aws-sdk-go-v2 v1.32.7 h1:ky5o35oENWi0JYWUZkB7WYvVPP+bcRF5/Iq7JWSb 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/config v1.28.7 h1:GduUnoTXlhkgnxTD93g1nv4tVPILbdNQOzav+Wpg7AE= -github.com/aws/aws-sdk-go-v2/config v1.28.7/go.mod h1:vZGX6GVkIE8uECSUHB6MWAUsd4ZcG2Yq/dMa4refR3M= 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/feature/ec2/imds v1.16.22 h1:kqOrpojG71DxJm/KDPO+Z/y1phm1JlC8/iT+5XRmAn8= -github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.22/go.mod h1:NtSFajXVVL8TA2QNngagVZmUtXciyrHOt7xgz4faS/M= 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/ini v1.8.1 h1:VaRN3TlFdd6KxX1x3ILT5ynH6HvKgqdiXoTxAF4HQcQ= -github.com/aws/aws-sdk-go-v2/internal/ini v1.8.1/go.mod h1:FbtygfRFze9usAadmnGJNc8KsP346kEe+y2/oyhGAGc= 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= @@ -34,12 +28,6 @@ github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.18.7 h1:Hi0KGbrnr57bEH 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/aws-sdk-go-v2/service/sso v1.24.8 h1:CvuUmnXI7ebaUAhbJcDy9YQx8wHR69eZ9I7q5hszt/g= -github.com/aws/aws-sdk-go-v2/service/sso v1.24.8/go.mod h1:XDeGv1opzwm8ubxddF0cgqkZWsyOtw4lr6dxwmb6YQg= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7 h1:F2rBfNAL5UyswqoeWv9zs74N/NanhK16ydHW1pahX6E= -github.com/aws/aws-sdk-go-v2/service/ssooidc v1.28.7/go.mod h1:JfyQ0g2JG8+Krq0EuZNnRwX0mU0HrwY/tG6JNfcqh4k= -github.com/aws/aws-sdk-go-v2/service/sts v1.33.3 h1:Xgv/hyNgvLda/M9l9qxXc4UFSgppnRczLxlMs5Ae/QY= -github.com/aws/aws-sdk-go-v2/service/sts v1.33.3/go.mod h1:5Gn+d+VaaRgsjewpMvGazt0WfcFO+Md4wLOuBfGR9Bc= 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= diff --git a/neo/load.go b/neo/load.go index 2e364b6c..3a185730 100644 --- a/neo/load.go +++ b/neo/load.go @@ -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 { diff --git a/neo/types.go b/neo/types.go index a692fb98..53c9ac24 100644 --- a/neo/types.go +++ b/neo/types.go @@ -7,6 +7,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" ) // DSL AI assistant @@ -18,6 +20,7 @@ type DSL struct { 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"` @@ -29,11 +32,18 @@ type DSL struct { 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:"-"` AssistantList []assistant.Assistant `json:"-" yaml:"-"` AssistantMaps map[string]assistant.Assistant `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 type Mention struct { ID string `json:"id"` From 55c0580fc38e64e2fe703cde72fb4d16da600e6e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 3 Jan 2025 16:37:43 +0800 Subject: [PATCH 4/5] Refactor Neo API assistant management and remove obsolete functions - Removed unused assistant-related functions and types from the Neo API, including local chat and file handling, to streamline the codebase. - Refactored the assistant initialization process, consolidating the loading and configuration of assistants, including Vision and Connector settings. - Updated the DSL struct to enhance clarity and maintainability, ensuring better configuration management for assistants. - Improved error handling and modularity in the assistant loading process, paving the way for future enhancements in AI functionalities. --- neo/assistant/api.go | 238 ++++++++++++++++++++++++++++++++- neo/assistant/assistant.go | 2 +- neo/assistant/load.go | 49 +++++++ neo/assistant/local/chat.go | 21 --- neo/assistant/local/file.go | 162 ---------------------- neo/assistant/local/local.go | 34 ----- neo/assistant/openai/chat.go | 118 ---------------- neo/assistant/openai/file.go | 167 ----------------------- neo/assistant/openai/openai.go | 51 ------- neo/assistant/openai/thread.go | 21 --- neo/assistant/types.go | 3 +- neo/hooks.go | 103 -------------- neo/load.go | 10 +- neo/neo.go | 184 ++----------------------- neo/types.go | 42 +++--- 15 files changed, 324 insertions(+), 881 deletions(-) delete mode 100644 neo/assistant/local/chat.go delete mode 100644 neo/assistant/local/file.go delete mode 100644 neo/assistant/local/local.go delete mode 100644 neo/assistant/openai/chat.go delete mode 100644 neo/assistant/openai/file.go delete mode 100644 neo/assistant/openai/openai.go delete mode 100644 neo/assistant/openai/thread.go diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 69c22a00..260ce88f 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -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,229 @@ 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{}{} + 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 +} diff --git a/neo/assistant/assistant.go b/neo/assistant/assistant.go index fd24eb9f..f4bc9fd0 100644 --- a/neo/assistant/assistant.go +++ b/neo/assistant/assistant.go @@ -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 diff --git a/neo/assistant/load.go b/neo/assistant/load.go index 734d612e..b17e8b02 100644 --- a/neo/assistant/load.go +++ b/neo/assistant/load.go @@ -12,6 +12,8 @@ import ( "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 +22,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 { @@ -73,6 +77,12 @@ func LoadBuiltIn() error { return err } + // Initialize the assistant + err = assistant.initialize() + if err != nil { + return err + } + sort++ loaded.Put(assistant) @@ -86,6 +96,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 +135,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 @@ -331,6 +356,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 +429,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 +} diff --git a/neo/assistant/local/chat.go b/neo/assistant/local/chat.go deleted file mode 100644 index 8c2518cb..00000000 --- a/neo/assistant/local/chat.go +++ /dev/null @@ -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 -} diff --git a/neo/assistant/local/file.go b/neo/assistant/local/file.go deleted file mode 100644 index cb7a692a..00000000 --- a/neo/assistant/local/file.go +++ /dev/null @@ -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 -} diff --git a/neo/assistant/local/local.go b/neo/assistant/local/local.go deleted file mode 100644 index 7a1386ff..00000000 --- a/neo/assistant/local/local.go +++ /dev/null @@ -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 -} diff --git a/neo/assistant/openai/chat.go b/neo/assistant/openai/chat.go deleted file mode 100644 index 1bf80a8c..00000000 --- a/neo/assistant/openai/chat.go +++ /dev/null @@ -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 -} diff --git a/neo/assistant/openai/file.go b/neo/assistant/openai/file.go deleted file mode 100644 index 062a688b..00000000 --- a/neo/assistant/openai/file.go +++ /dev/null @@ -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 -} diff --git a/neo/assistant/openai/openai.go b/neo/assistant/openai/openai.go deleted file mode 100644 index 38f5135c..00000000 --- a/neo/assistant/openai/openai.go +++ /dev/null @@ -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() {} diff --git a/neo/assistant/openai/thread.go b/neo/assistant/openai/thread.go deleted file mode 100644 index 8f5975e8..00000000 --- a/neo/assistant/openai/thread.go +++ /dev/null @@ -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() {} diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 5cec9ba6..480f8146 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -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 diff --git a/neo/hooks.go b/neo/hooks.go index 775f492f..4ec7dabb 100644 --- a/neo/hooks.go +++ b/neo/hooks.go @@ -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 -} diff --git a/neo/load.go b/neo/load.go index 3a185730..4ee561e4 100644 --- a/neo/load.go +++ b/neo/load.go @@ -155,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 { @@ -167,7 +175,7 @@ func (neo *DSL) initAssistant() error { return err } - Neo.Assistant = defaultAssistant.API + Neo.Assistant = defaultAssistant return nil } diff --git a/neo/neo.go b/neo/neo.go index f072dc1c..b114b0d8 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -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 { diff --git a/neo/types.go b/neo/types.go index 53c9ac24..e3511168 100644 --- a/neo/types.go +++ b/neo/types.go @@ -13,29 +13,25 @@ import ( // 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"` - 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"` - 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:"-"` - Vision *vision.Vision `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 From 0d676b33486cf3b92e7557cb2bd67dd4334da078 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 3 Jan 2025 17:17:04 +0800 Subject: [PATCH 5/5] Update Go module dependencies and enhance assistant message handling in Neo API - Added the github.com/spf13/cast dependency to improve type conversion capabilities. - Enhanced the requestMessages function to include prompts, allowing for more dynamic message handling. - Updated the LoadBuiltIn function to ensure the "Built-in" tag is added if it doesn't already exist, improving assistant organization. - Refactored sorting logic to utilize the cast package for better type safety. These changes support improved functionality and maintainability in the assistant management process. --- go.mod | 1 + go.sum | 3 +++ neo/assistant/api.go | 19 +++++++++++++++++++ neo/assistant/load.go | 23 ++++++++++++++++++++--- 4 files changed, 43 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 082cb348..d018c183 100644 --- a/go.mod +++ b/go.mod @@ -102,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 diff --git a/go.sum b/go.sum index 09a2c82a..2efec00f 100644 --- a/go.sum +++ b/go.sum @@ -160,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= @@ -219,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= diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 260ce88f..8cad3890 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -82,6 +82,25 @@ func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{ 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) diff --git a/neo/assistant/load.go b/neo/assistant/load.go index b17e8b02..99cfb2be 100644 --- a/neo/assistant/load.go +++ b/neo/assistant/load.go @@ -8,6 +8,7 @@ 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" @@ -66,11 +67,27 @@ 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 { @@ -284,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