From d4d824a22c313c9776260b9e5693174040de5ba3 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Mon, 6 Jan 2025 11:27:44 +0800
Subject: [PATCH] Refactor Neo API upload functionality and enhance environment
 variable handling

- Updated the Upload method in the DSL to default to the assistant in context, improving error handling for missing assistant IDs.
- Introduced parseEnvValue and convertOptions functions in the Vision module to support parsing environment variables in configuration options, enhancing flexibility.
- Refactored storage and model initialization in the Vision service to utilize converted options, ensuring environment variables are correctly applied.

These changes improve the robustness and maintainability of the Neo API, paving the way for better configuration management and error handling in assistant functionalities.
---
 neo/neo.go           | 19 ++++++++++---------
 neo/vision/vision.go | 39 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/neo/neo.go b/neo/neo.go
index b114b0d8..f9d56b35 100644
--- a/neo/neo.go
+++ b/neo/neo.go
@@ -219,15 +219,16 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
 		Option:      option,
 	}
 
-	res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c)
-	if err != nil {
-		return nil, err
-	}
-
-	// Select Assistant
-	ast, err := neo.Select(res.AssistantID)
-	if err != nil {
-		return nil, err
+	// Default use the assistant in context
+	ast := neo.Assistant
+	if ctx.ChatID == "" {
+		if ctx.AssistantID == "" {
+			return nil, fmt.Errorf("assistant_id is required")
+		}
+		ast, err = neo.Select(ctx.AssistantID)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	return ast.Upload(ctx, tmpfile, reader, option)
diff --git a/neo/vision/vision.go b/neo/vision/vision.go
index 79c5f372..0886a900 100644
--- a/neo/vision/vision.go
+++ b/neo/vision/vision.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"os"
 	"strings"
 	"time"
 
@@ -13,6 +14,30 @@ import (
 	"github.com/yaoapp/yao/neo/vision/driver/s3"
 )
 
+// parseEnvValue parse environment variable if the value starts with $ENV.
+func parseEnvValue(value string) string {
+	if strings.HasPrefix(value, "$ENV.") {
+		envKey := strings.TrimPrefix(value, "$ENV.")
+		if envVal := os.Getenv(envKey); envVal != "" {
+			return envVal
+		}
+	}
+	return value
+}
+
+// convertOptions convert interface{} options map to string map and parse environment variables
+func convertOptions(options map[string]interface{}) map[string]interface{} {
+	converted := make(map[string]interface{})
+	for k, v := range options {
+		if str, ok := v.(string); ok {
+			converted[k] = parseEnvValue(str)
+		} else {
+			converted[k] = v
+		}
+	}
+	return converted
+}
+
 // Vision the vision service
 type Vision struct {
 	storage driver.Storage
@@ -22,20 +47,24 @@ type Vision struct {
 // New create a new vision service
 func New(cfg *driver.Config) (*Vision, error) {
 
+	// Parse environment variables in options
+	storageOptions := convertOptions(cfg.Storage.Options)
+	modelOptions := convertOptions(cfg.Model.Options)
+
 	// Create storage driver
 	var storage driver.Storage
 	var err error
 	switch cfg.Storage.Driver {
 	case "local":
-		storage, err = local.New(cfg.Storage.Options)
+		storage, err = local.New(storageOptions)
 	case "s3":
 		// Convert expiration string to duration if present
-		if exp, ok := cfg.Storage.Options["expiration"].(string); ok {
+		if exp, ok := storageOptions["expiration"].(string); ok {
 			if duration, err := time.ParseDuration(exp); err == nil {
-				cfg.Storage.Options["expiration"] = duration
+				storageOptions["expiration"] = duration
 			}
 		}
-		storage, err = s3.New(cfg.Storage.Options)
+		storage, err = s3.New(storageOptions)
 	default:
 		return nil, fmt.Errorf("storage driver %s not supported", cfg.Storage.Driver)
 	}
@@ -47,7 +76,7 @@ func New(cfg *driver.Config) (*Vision, error) {
 	var model driver.Model
 	switch cfg.Model.Driver {
 	case "openai":
-		model, err = openai.New(cfg.Model.Options)
+		model, err = openai.New(modelOptions)
 	default:
 		return nil, fmt.Errorf("model driver %s not supported", cfg.Model.Driver)
 	}