5f7b2e866c
- Removed obsolete file upload/download functions and related constants from the Assistant struct, streamlining the codebase. - Updated the connector type in load tests to use "gpt-3_5-turbo" for improved compatibility with current models. - Introduced VisionCapableModels to support a range of models with vision capabilities, enhancing AI functionality. - Enhanced the File struct to include additional fields for description, URL, and document IDs, improving data management for file handling. These changes improve the overall structure and maintainability of the Neo API, paving the way for future enhancements in assistant capabilities.
338 lines
8.8 KiB
Go
338 lines
8.8 KiB
Go
package assistant
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"io"
|
|
"mime/multipart"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/yaoapp/gou/fs"
|
|
gourag "github.com/yaoapp/gou/rag"
|
|
"github.com/yaoapp/gou/rag/driver"
|
|
"github.com/yaoapp/yao/config"
|
|
neovision "github.com/yaoapp/yao/neo/vision"
|
|
vdriver "github.com/yaoapp/yao/neo/vision/driver"
|
|
"github.com/yaoapp/yao/test"
|
|
)
|
|
|
|
var (
|
|
// 1x1 transparent PNG for testing
|
|
testImageBase64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
|
)
|
|
|
|
func TestUpload(t *testing.T) {
|
|
test.Prepare(t, config.Conf)
|
|
defer test.Clean()
|
|
|
|
ast := setupTestAssistant(t)
|
|
ctx := context.Background()
|
|
|
|
t.Run("Basic File Upload", func(t *testing.T) {
|
|
content := []byte("test content")
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.txt",
|
|
Size: int64(len(content)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "text/plain")
|
|
|
|
reader := bytes.NewReader(content)
|
|
fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
|
|
"sid": "test-user",
|
|
"chat_id": "test-chat",
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, fileResp)
|
|
assert.Contains(t, fileResp.ID, "test-assistant/test-user/test-chat")
|
|
assert.Equal(t, len(content), fileResp.Bytes)
|
|
assert.Equal(t, "text/plain", fileResp.ContentType)
|
|
})
|
|
|
|
t.Run("File Size Limit", func(t *testing.T) {
|
|
content := make([]byte, MaxSize+1)
|
|
file := &multipart.FileHeader{
|
|
Filename: "large.txt",
|
|
Size: int64(len(content)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "text/plain")
|
|
|
|
reader := bytes.NewReader(content)
|
|
_, err := ast.Upload(ctx, file, reader, nil)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "exceeds the maximum size")
|
|
})
|
|
|
|
t.Run("Invalid Content Type", func(t *testing.T) {
|
|
content := []byte("test")
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.invalid",
|
|
Size: int64(len(content)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "invalid/type")
|
|
|
|
reader := bytes.NewReader(content)
|
|
_, err := ast.Upload(ctx, file, reader, nil)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "not allowed")
|
|
})
|
|
}
|
|
|
|
func TestUploadWithRAG(t *testing.T) {
|
|
test.Prepare(t, config.Conf)
|
|
defer test.Clean()
|
|
|
|
ast := setupTestAssistant(t)
|
|
ragEngine, ragUploader, ragVectorizer := setupTestRAG(t)
|
|
SetRAG(ragEngine, ragUploader, ragVectorizer, RAGSetting{IndexPrefix: "test_"})
|
|
defer func() {
|
|
rag = nil // Completely reset the global rag variable
|
|
}()
|
|
ctx := context.Background()
|
|
|
|
t.Run("Text File with RAG", func(t *testing.T) {
|
|
content := []byte("This is a test document for RAG indexing")
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.txt",
|
|
Size: int64(len(content)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "text/plain")
|
|
|
|
reader := bytes.NewReader(content)
|
|
fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
|
|
"sid": "test-user",
|
|
"chat_id": "test-chat",
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, fileResp)
|
|
assert.NotEmpty(t, fileResp.DocIDs, "Document IDs should not be empty")
|
|
|
|
// Wait for indexing to complete
|
|
time.Sleep(500 * time.Millisecond)
|
|
|
|
// Verify the file was indexed by checking if it exists in RAG
|
|
exists, err := ragEngine.HasDocument(ctx, "test_test-assistant-test-user-test-chat", fileResp.DocIDs[0])
|
|
assert.NoError(t, err)
|
|
assert.True(t, exists, "Document should exist in RAG index")
|
|
})
|
|
|
|
t.Run("Non-Text File with RAG", func(t *testing.T) {
|
|
imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.png",
|
|
Size: int64(len(imgData)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "image/png")
|
|
|
|
reader := bytes.NewReader(imgData)
|
|
fileResp, err := ast.Upload(ctx, file, reader, nil)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, fileResp)
|
|
// Verify the file was not indexed
|
|
exists, err := ragEngine.HasDocument(ctx, "test_test-assistant", fileResp.ID)
|
|
assert.NoError(t, err)
|
|
assert.False(t, exists)
|
|
})
|
|
}
|
|
|
|
func setupTestAssistant(t *testing.T) *Assistant {
|
|
ast := &Assistant{
|
|
ID: "test-assistant",
|
|
Name: "Test Assistant",
|
|
Connector: "test-connector",
|
|
}
|
|
return ast
|
|
}
|
|
|
|
func setupTestRAG(t *testing.T) (driver.Engine, driver.FileUpload, driver.Vectorizer) {
|
|
// Get test config
|
|
openaiKey := os.Getenv("OPENAI_API_KEY")
|
|
if openaiKey == "" {
|
|
t.Skip("OPENAI_API_KEY not set")
|
|
}
|
|
|
|
vectorizeConfig := driver.VectorizeConfig{
|
|
Model: os.Getenv("VECTORIZER_MODEL"),
|
|
Options: map[string]string{
|
|
"api_key": openaiKey,
|
|
},
|
|
}
|
|
|
|
// Qdrant config
|
|
host := os.Getenv("QDRANT_HOST")
|
|
if host == "" {
|
|
host = "localhost"
|
|
}
|
|
|
|
port := os.Getenv("QDRANT_PORT")
|
|
if port == "" {
|
|
port = "6334"
|
|
}
|
|
|
|
// Create vectorizer
|
|
vectorizer, err := gourag.NewVectorizer(gourag.DriverOpenAI, vectorizeConfig)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// Create engine
|
|
engine, err := gourag.NewEngine(gourag.DriverQdrant, driver.IndexConfig{
|
|
Options: map[string]string{
|
|
"host": host,
|
|
"port": port,
|
|
"api_key": "",
|
|
},
|
|
}, vectorizer)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// Create file upload
|
|
fileUpload, err := gourag.NewFileUpload(gourag.DriverQdrant, engine, vectorizer)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
return engine, fileUpload, vectorizer
|
|
}
|
|
|
|
func setupTestVision(t *testing.T) *neovision.Vision {
|
|
// Create test data directory
|
|
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)
|
|
|
|
cfg := &vdriver.Config{
|
|
Storage: vdriver.StorageConfig{
|
|
Driver: "local",
|
|
Options: map[string]interface{}{
|
|
"path": "/__vision_test",
|
|
"compression": true,
|
|
},
|
|
},
|
|
Model: vdriver.ModelConfig{
|
|
Driver: "openai",
|
|
Options: map[string]interface{}{
|
|
"api_key": os.Getenv("OPENAI_API_KEY"),
|
|
"model": os.Getenv("VISION_MODEL"),
|
|
},
|
|
},
|
|
}
|
|
|
|
v, err := neovision.New(cfg)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return v
|
|
}
|
|
|
|
func TestUploadWithVision(t *testing.T) {
|
|
test.Prepare(t, config.Conf)
|
|
defer test.Clean()
|
|
|
|
ast := setupTestAssistant(t)
|
|
vision := setupTestVision(t)
|
|
SetVision(vision)
|
|
defer func() {
|
|
vision = nil // Completely reset the global vision variable
|
|
}()
|
|
ctx := context.Background()
|
|
|
|
t.Run("Image with Vision-Capable Model", func(t *testing.T) {
|
|
imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.png",
|
|
Size: int64(len(imgData)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "image/png")
|
|
|
|
reader := bytes.NewReader(imgData)
|
|
fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
|
|
"model": "gpt-4-vision-preview",
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, fileResp)
|
|
assert.NotEmpty(t, fileResp.URL)
|
|
assert.Empty(t, fileResp.Description)
|
|
})
|
|
|
|
t.Run("Image with Non-Vision Model", func(t *testing.T) {
|
|
imgData, _ := base64.StdEncoding.DecodeString(testImageBase64)
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.png",
|
|
Size: int64(len(imgData)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "image/png")
|
|
|
|
reader := bytes.NewReader(imgData)
|
|
fileResp, err := ast.Upload(ctx, file, reader, map[string]interface{}{
|
|
"model": "gpt-4",
|
|
"vision_prompt": "What's in this image?",
|
|
})
|
|
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, fileResp)
|
|
assert.Empty(t, fileResp.URL)
|
|
assert.NotEmpty(t, fileResp.Description)
|
|
})
|
|
}
|
|
|
|
func TestDownload(t *testing.T) {
|
|
test.Prepare(t, config.Conf)
|
|
defer test.Clean()
|
|
|
|
ast := setupTestAssistant(t)
|
|
ctx := context.Background()
|
|
|
|
t.Run("Download Existing File", func(t *testing.T) {
|
|
// First upload a file
|
|
content := []byte("test content")
|
|
file := &multipart.FileHeader{
|
|
Filename: "test.txt",
|
|
Size: int64(len(content)),
|
|
}
|
|
file.Header = make(map[string][]string)
|
|
file.Header.Set("Content-Type", "text/plain")
|
|
|
|
reader := bytes.NewReader(content)
|
|
fileResp, err := ast.Upload(ctx, file, reader, nil)
|
|
assert.NoError(t, err)
|
|
|
|
// Then download it
|
|
downloadResp, err := ast.Download(ctx, fileResp.ID)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, downloadResp)
|
|
assert.True(t, strings.HasPrefix(downloadResp.ContentType, "text/plain"), "Content-Type should start with text/plain")
|
|
assert.Equal(t, ".txt", downloadResp.Extension)
|
|
|
|
// Verify content
|
|
downloaded, err := io.ReadAll(downloadResp.Reader)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, content, downloaded)
|
|
})
|
|
|
|
t.Run("Download Non-Existent File", func(t *testing.T) {
|
|
_, err := ast.Download(ctx, "non-existent-file")
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "not found")
|
|
})
|
|
}
|