Refactor Neo API assistant functionality and enhance model capabilities
- 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.
This commit is contained in:
Max
parent
0d676b3348
commit
5f7b2e866c
5 changed files with 685 additions and 116 deletions
|
|
@ -2,14 +2,9 @@ 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"
|
||||
|
|
@ -45,22 +40,6 @@ func GetByConnector(connector string, name string) (*Assistant, error) {
|
|||
return assistant, nil
|
||||
}
|
||||
|
||||
// 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 {
|
||||
|
|
@ -175,94 +154,6 @@ func (ast *Assistant) withAttachments(ctx context.Context, msg *chatMessage.Mess
|
|||
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")
|
||||
|
|
|
|||
312
neo/assistant/attachment.go
Normal file
312
neo/assistant/attachment.go
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
package assistant
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
"github.com/yaoapp/gou/fs"
|
||||
"github.com/yaoapp/gou/rag/driver"
|
||||
)
|
||||
|
||||
// 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 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)
|
||||
}
|
||||
|
||||
// Get chat ID and session ID from options
|
||||
chatID := ""
|
||||
sid := ""
|
||||
if v, ok := option["chat_id"].(string); ok {
|
||||
chatID = v
|
||||
}
|
||||
if v, ok := option["sid"].(string); ok {
|
||||
sid = v
|
||||
}
|
||||
|
||||
// Generate file ID with namespace
|
||||
fileID, err := ast.generateFileID(file.Filename, sid, chatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Upload file to storage
|
||||
data, err := fs.Get("data")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
_, err = data.Write(fileID, reader, 0644)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Create file response
|
||||
fileResp := &File{
|
||||
ID: fileID,
|
||||
Filename: fileID,
|
||||
ContentType: contentType,
|
||||
Bytes: int(file.Size),
|
||||
CreatedAt: int(time.Now().Unix()),
|
||||
}
|
||||
|
||||
// Handle RAG if available
|
||||
if err := ast.handleRAG(ctx, fileResp, reader); err != nil {
|
||||
return nil, fmt.Errorf("RAG handling error: %s", err.Error())
|
||||
}
|
||||
|
||||
// Handle Vision if available
|
||||
if err := ast.handleVision(ctx, fileResp, option); err != nil {
|
||||
return nil, fmt.Errorf("Vision handling error: %s", err.Error())
|
||||
}
|
||||
|
||||
return fileResp, nil
|
||||
}
|
||||
|
||||
// generateFileID generates a file ID with proper namespace
|
||||
func (ast *Assistant) generateFileID(filename string, sid string, chatID string) (string, error) {
|
||||
ext := filepath.Ext(filename)
|
||||
hash := fmt.Sprintf("%x", sha256.Sum256([]byte(filename)))[:8]
|
||||
date := time.Now().Format("20060102")
|
||||
|
||||
// Build namespace
|
||||
namespace := fmt.Sprintf("__assistants/%s", ast.ID)
|
||||
if sid != "" {
|
||||
namespace = fmt.Sprintf("%s/%s", namespace, sid)
|
||||
if chatID != "" {
|
||||
namespace = fmt.Sprintf("%s/%s", namespace, chatID)
|
||||
}
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s/%s/%s%s", namespace, date, hash, ext), nil
|
||||
}
|
||||
|
||||
// handleRAG handles the file with RAG if available
|
||||
func (ast *Assistant) handleRAG(ctx context.Context, file *File, reader io.Reader) error {
|
||||
if rag == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Only handle text-based files
|
||||
if !strings.HasPrefix(file.ContentType, "text/") {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Reset reader to beginning
|
||||
if seeker, ok := reader.(io.Seeker); ok {
|
||||
if _, err := seeker.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Extract sid and chat_id from file path
|
||||
parts := strings.Split(file.ID, "/")
|
||||
indexName := fmt.Sprintf("%s%s", rag.Setting.IndexPrefix, ast.ID) // Default: prefix-assistant
|
||||
|
||||
if len(parts) >= 4 { // Has sid
|
||||
sid := parts[2]
|
||||
indexName = fmt.Sprintf("%s%s-%s", rag.Setting.IndexPrefix, ast.ID, sid) // prefix-assistant-user
|
||||
|
||||
if len(parts) >= 5 { // Has chat_id
|
||||
chatID := parts[3]
|
||||
indexName = fmt.Sprintf("%s%s-%s-%s", rag.Setting.IndexPrefix, ast.ID, sid, chatID) // prefix-assistant-user-chat
|
||||
}
|
||||
}
|
||||
|
||||
// Check if index exists
|
||||
exists, err := rag.Engine.HasIndex(ctx, indexName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("check index error: %s", err.Error())
|
||||
}
|
||||
|
||||
// Create index if not exists
|
||||
if !exists {
|
||||
err = rag.Engine.CreateIndex(ctx, driver.IndexConfig{Name: indexName})
|
||||
if err != nil {
|
||||
return fmt.Errorf("create index error: %s", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// Reset reader again after checking index
|
||||
if seeker, ok := reader.(io.Seeker); ok {
|
||||
if _, err := seeker.Seek(0, io.SeekStart); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Upload and index the file
|
||||
result, err := rag.Uploader.Upload(ctx, reader, driver.FileUploadOptions{
|
||||
Async: false,
|
||||
ChunkSize: 1024, // Default chunk size
|
||||
ChunkOverlap: 256, // Default overlap
|
||||
IndexName: indexName,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
return fmt.Errorf("upload error: %s", err.Error())
|
||||
}
|
||||
|
||||
if len(result.Documents) == 0 {
|
||||
return fmt.Errorf("no documents indexed")
|
||||
}
|
||||
|
||||
// Store the document IDs
|
||||
docIDs := make([]string, len(result.Documents))
|
||||
for i, doc := range result.Documents {
|
||||
docIDs[i] = doc.DocID
|
||||
}
|
||||
file.DocIDs = docIDs
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleVision handles the file with Vision if available
|
||||
func (ast *Assistant) handleVision(ctx context.Context, file *File, option map[string]interface{}) error {
|
||||
if vision == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check if file is an image
|
||||
if !strings.HasPrefix(file.ContentType, "image/") {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get model from options
|
||||
model := ""
|
||||
if v, ok := option["model"].(string); ok {
|
||||
model = v
|
||||
}
|
||||
|
||||
// Reset reader for vision service
|
||||
data, err := fs.Get("data")
|
||||
if err != nil {
|
||||
return fmt.Errorf("get filesystem error: %s", err.Error())
|
||||
}
|
||||
|
||||
exists, err := data.Exists(file.ID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("check file error: %s", err.Error())
|
||||
}
|
||||
if !exists {
|
||||
return fmt.Errorf("file %s not found", file.ID)
|
||||
}
|
||||
|
||||
// Read file content into memory
|
||||
imgData, err := data.ReadFile(file.ID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read file error: %s", err.Error())
|
||||
}
|
||||
|
||||
if VisionCapableModels[model] {
|
||||
// For vision-capable models, upload to vision service to get URL
|
||||
resp, err := vision.Upload(ctx, file.Filename, bytes.NewReader(imgData), file.ContentType)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vision upload error: %s", err.Error())
|
||||
}
|
||||
file.URL = resp.URL // Store the URL for vision-capable models to use
|
||||
} else {
|
||||
// For non-vision models, get image description
|
||||
prompt := "Describe this image in detail."
|
||||
if v, ok := option["vision_prompt"].(string); ok {
|
||||
prompt = v
|
||||
}
|
||||
|
||||
// Upload to vision service first Compress image
|
||||
resp, err := vision.Upload(ctx, file.Filename, bytes.NewReader(imgData), file.ContentType)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vision upload error: %s", err.Error())
|
||||
}
|
||||
|
||||
// Analyze using base64 data
|
||||
result, err := vision.Analyze(ctx, resp.FileID, prompt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("vision analyze error: %s", err.Error())
|
||||
}
|
||||
|
||||
// Extract description text from response
|
||||
if desc, ok := result.Description["text"].(string); ok {
|
||||
file.Description = desc
|
||||
} else {
|
||||
// Convert the entire description to JSON string as fallback
|
||||
bytes, err := jsoniter.Marshal(result.Description)
|
||||
if err == nil {
|
||||
file.Description = string(bytes)
|
||||
}
|
||||
}
|
||||
}
|
||||
return 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
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
338
neo/assistant/attachment_test.go
Normal file
338
neo/assistant/attachment_test.go
Normal file
|
|
@ -0,0 +1,338 @@
|
|||
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")
|
||||
})
|
||||
}
|
||||
|
|
@ -53,7 +53,7 @@ func TestLoad_LoadStore(t *testing.T) {
|
|||
"assistant_id": "test-id",
|
||||
"name": "Test Assistant",
|
||||
"avatar": "test-avatar",
|
||||
"connector": "test-connector",
|
||||
"connector": "gpt-3_5-turbo",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
|
@ -67,7 +67,7 @@ func TestLoad_LoadStore(t *testing.T) {
|
|||
assert.Equal(t, "test-id", assistant.ID)
|
||||
assert.Equal(t, "Test Assistant", assistant.Name)
|
||||
assert.Equal(t, "test-avatar", assistant.Avatar)
|
||||
assert.Equal(t, "test-connector", assistant.Connector)
|
||||
assert.Equal(t, "gpt-3_5-turbo", assistant.Connector)
|
||||
|
||||
// Test cache functionality
|
||||
assistant2, err := LoadStore("test-id")
|
||||
|
|
|
|||
|
|
@ -70,13 +70,41 @@ type Assistant struct {
|
|||
openai *api.OpenAI // OpenAI API
|
||||
}
|
||||
|
||||
// VisionCapableModels list of LLM models that support vision capabilities
|
||||
var VisionCapableModels = map[string]bool{
|
||||
// OpenAI Models
|
||||
"gpt-4-vision-preview": true,
|
||||
"gpt-4v": true, // Alias for gpt-4-vision-preview
|
||||
|
||||
// Anthropic Models
|
||||
"claude-3-opus": true, // Most capable Claude model
|
||||
"claude-3-sonnet": true, // Balanced Claude model
|
||||
"claude-3-haiku": true, // Fast and efficient Claude model
|
||||
|
||||
// Google Models
|
||||
"gemini-pro-vision": true,
|
||||
|
||||
// Open Source Models
|
||||
"llava-13b": true,
|
||||
"cogvlm": true,
|
||||
"qwen-vl": true,
|
||||
"yi-vl": true,
|
||||
|
||||
// Custom Models
|
||||
"gpt-4o": true, // Custom OpenAI compatible model
|
||||
"gpt-4o-mini": true, // Custom OpenAI compatible model - mini version
|
||||
}
|
||||
|
||||
// File the file
|
||||
type File struct {
|
||||
ID string `json:"file_id"`
|
||||
Bytes int `json:"bytes"`
|
||||
CreatedAt int `json:"created_at"`
|
||||
Filename string `json:"filename"`
|
||||
ContentType string `json:"content_type"`
|
||||
ID string `json:"file_id"`
|
||||
Bytes int `json:"bytes"`
|
||||
CreatedAt int `json:"created_at"`
|
||||
Filename string `json:"filename"`
|
||||
ContentType string `json:"content_type"`
|
||||
Description string `json:"description,omitempty"` // Vision analysis result or other description
|
||||
URL string `json:"url,omitempty"` // Vision URL for vision-capable models
|
||||
DocIDs []string `json:"doc_ids,omitempty"` // RAG document IDs
|
||||
}
|
||||
|
||||
// FileResponse represents a file download response
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue