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

Merge pull request #816 from trheyi/main

Browse source 
Refactor Neo API and Enhance Assistant Capabilities (Dev)
This commit is contained in:
Max 2025-01-13 18:20:22 +08:00 • committed by GitHub
commit 9c7b9cf9a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 1653 additions and 388 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
neo/api.go
View file

@ -15,6 +15,7 @@ import (
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/helper"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
"github.com/yaoapp/yao/neo/store"
)
@ -173,7 +174,7 @@ func (neo *DSL) handleUpload(c *gin.Context) {
}
// Set the context
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
// Upload the file
@ -214,7 +215,7 @@ func (neo *DSL) handleChat(c *gin.Context) {
}
// Set the context with validated chat_id
ctx, cancel := NewContextWithCancel(sid, chatID, c.Query("context"))
ctx, cancel := chatctx.NewWithCancel(sid, chatID, c.Query("context"))
defer cancel()
neo.Answer(ctx, content, c)
@ -297,7 +298,7 @@ func (neo *DSL) handleDownload(c *gin.Context) {
}
// Set the context
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
// Download the file
@ -537,7 +538,7 @@ func (neo *DSL) handleChatUpdate(c *gin.Context) {
// If content is not empty, Generate the chat title
if body.Content != "" {
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
defer cancel()
title, err := neo.GenerateChatTitle(ctx, body.Content, c, true)
@ -729,7 +730,7 @@ func (neo *DSL) handleGenerateTitle(c *gin.Context) {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
@ -780,7 +781,7 @@ func (neo *DSL) handleGeneratePrompts(c *gin.Context) {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE
@ -831,7 +832,7 @@ func (neo *DSL) handleGenerateCustom(c *gin.Context) {
return
}
ctx, cancel := NewContextWithCancel(resp.sid, c.Query("chat_id"), "")
ctx, cancel := chatctx.NewWithCancel(resp.sid, c.Query("chat_id"), "")
defer cancel()
// Use silent mode for regular HTTP requests, streaming for SSE

357
neo/assistant/api.go
View file

@ -2,16 +2,14 @@ package assistant
import (
"context"
"crypto/sha256"
"encoding/base64"
"fmt"
"io"
"mime/multipart"
"path/filepath"
"strings"
"time"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/fs"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
chatMessage "github.com/yaoapp/yao/neo/message"
)
@ -45,24 +43,238 @@ 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",
// Execute implements the execute functionality
func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error {
messages, err := ast.withHistory(ctx, input)
if err != nil {
return err
}
options = ast.withOptions(options)
// Run init hook
res, err := ast.HookInit(c, ctx, messages, options)
if err != nil {
return err
}
// Switch to the new assistant if necessary
if res.AssistantID != ctx.AssistantID {
newAst, err := Get(res.AssistantID)
if err != nil {
return err
}
*ast = *newAst
}
// Handle next action
if res.Next != nil {
switch res.Next.Action {
case "exit":
return nil
// Add other actions here if needed
}
}
// Update options if provided
if res.Options != nil {
options = res.Options
}
// messages
if res.Input != nil {
messages = res.Input
}
// Only proceed with chat stream if no specific next action was handled
return ast.handleChatStream(c, ctx, messages, options)
}
// MaxSize 20M max file size
var MaxSize int64 = 20 * 1024 * 1024
// handleChatStream manages the streaming chat interaction with the AI
func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{}) error {
clientBreak := make(chan bool, 1)
done := make(chan bool, 1)
content := []byte{}
// Chat with AI in background
go func() {
err := ast.streamChat(c, ctx, messages, options, clientBreak, done, &content)
if err != nil {
chatMessage.New().Error(err).Done().Write(c.Writer)
}
ast.saveChatHistory(ctx, messages, content)
done <- true
}()
// Wait for completion or client disconnect
select {
case <-done:
return nil
case <-c.Writer.CloseNotify():
clientBreak <- true
return nil
}
}
// streamChat handles the streaming chat interaction
func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{},
clientBreak chan bool, done chan bool, content *[]byte) error {
return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
select {
case <-clientBreak:
return 0 // break
default:
msg := chatMessage.NewOpenAI(data)
if msg == nil {
return 1 // continue
}
// Handle error
if msg.Type == "error" {
value := msg.String()
res, hookErr := ast.HookFail(c, ctx, messages, string(*content), fmt.Errorf("%s", value))
if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
value = res.Output
if res.Error != "" {
value = res.Error
}
}
chatMessage.New().Error(value).Done().Write(c.Writer)
return 0 // break
}
// Append content and send message
*content = msg.Append(*content)
value := msg.String()
if value != "" {
// Handle stream
res, err := ast.HookStream(c, ctx, messages, string(*content))
if err == nil && res != nil {
if res.Output != "" {
value = res.Output
}
if res.Next != nil && res.Next.Action == "exit" {
done <- true
return 0 // break
}
if res.Silent {
return 1 // continue
}
}
chatMessage.New().
Map(map[string]interface{}{
"text": value,
"done": msg.IsDone,
}).
Write(c.Writer)
}
// Complete the stream
if msg.IsDone {
// if value == "" {
// msg.Write(c.Writer)
// }
// Call HookDone
res, hookErr := ast.HookDone(c, ctx, messages, string(*content))
if hookErr == nil && res != nil {
if res.Output != "" {
chatMessage.New().
Map(map[string]interface{}{
"text": res.Output,
"done": true,
}).
Write(c.Writer)
}
if res.Next != nil && res.Next.Action == "exit" {
done <- true
return 0 // break
}
} else if value != "" {
chatMessage.New().
Map(map[string]interface{}{
"text": value,
"done": true,
}).
Write(c.Writer)
}
done <- true
return 0 // break
}
return 1 // continue
}
})
}
// saveChatHistory saves the chat history if storage is available
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content []byte) {
if len(content) > 0 && ctx.Sid != "" && len(messages) > 0 {
storage.SaveHistory(
ctx.Sid,
[]map[string]interface{}{
{"role": "user", "content": messages[len(messages)-1].Content(), "name": ctx.Sid},
{"role": "assistant", "content": string(content), "name": ctx.Sid},
},
ctx.ChatID,
nil,
)
}
}
func (ast *Assistant) withOptions(options map[string]interface{}) map[string]interface{} {
if options == nil {
options = map[string]interface{}{}
}
if ast.Options != nil {
for key, value := range ast.Options {
options[key] = value
}
}
return options
}
func (ast *Assistant) withPrompts(messages []message.Message) []message.Message {
if ast.Prompts != nil {
for _, prompt := range ast.Prompts {
name := ast.Name
if prompt.Name != "" {
name = prompt.Name
}
messages = append(messages, *message.New().Map(map[string]interface{}{"role": prompt.Role, "content": prompt.Content, "name": name}))
}
}
return messages
}
func (ast *Assistant) withHistory(ctx chatctx.Context, input string) ([]message.Message, error) {
messages := []message.Message{}
messages = ast.withPrompts(messages)
if storage != nil {
history, err := storage.GetHistory(ctx.Sid, ctx.ChatID)
if err != nil {
return nil, err
}
// Add history messages
for _, h := range history {
messages = append(messages, *message.New().Map(h))
}
}
// Add user message
messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": input, "name": ctx.Sid}))
return messages, nil
}
// 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 {
func (ast *Assistant) Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error {
if ast.openai == nil {
return fmt.Errorf("openai is not initialized")
}
@ -80,13 +292,12 @@ func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{
return nil
}
func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
func (ast *Assistant) requestMessages(ctx context.Context, messages []message.Message) ([]map[string]interface{}, error) {
newMessages := []map[string]interface{}{}
// With Prompts
if ast.Prompts != nil {
for _, prompt := range ast.Prompts {
message := map[string]interface{}{
msg := map[string]interface{}{
"role": prompt.Role,
"content": prompt.Content,
}
@ -96,20 +307,20 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string
name = prompt.Name
}
message["name"] = name
newMessages = append(newMessages, message)
msg["name"] = name
newMessages = append(newMessages, msg)
}
}
length := len(messages)
for index, message := range messages {
role, ok := message["role"].(string)
if !ok {
role := message.Role
if role == "" {
return nil, fmt.Errorf("role must be string")
}
content, ok := message["content"].(string)
if !ok {
content := message.Text
if content == "" {
return nil, fmt.Errorf("content must be string")
}
@ -118,7 +329,7 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string
"content": content,
}
if name, ok := message["name"].(string); ok {
if name := message.Name; name != "" {
newMessage["name"] = name
}
@ -175,94 +386,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")

322
neo/assistant/attachment.go Normal file
View file

@ -0,0 +1,322 @@
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, option); 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, option map[string]interface{}) error {
if rag == nil {
return nil
}
// Check if RAG processing is enabled
if option, ok := option["rag"].(bool); !ok || !option {
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 Vision processing is enabled
if option, ok := option["vision"].(bool); !ok || !option {
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
}

363
neo/assistant/attachment_test.go Normal file
View file

@ -0,0 +1,363 @@
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()
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()
ragEngine, ragUploader, ragVectorizer := setupTestRAG(t)
SetRAG(ragEngine, ragUploader, ragVectorizer, RAGSetting{IndexPrefix: "test_"})
defer func() {
rag = nil
}()
ctx := context.Background()
t.Run("Text File with RAG Enabled", 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",
"rag": true,
})
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
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("Text File with RAG Disabled", func(t *testing.T) {
content := []byte("This is a test document with RAG disabled")
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",
"rag": false,
})
assert.NoError(t, err)
assert.NotNil(t, fileResp)
assert.Empty(t, fileResp.DocIDs, "Document IDs should be empty when RAG is disabled")
})
}
func setupTestAssistant() *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()
vision := setupTestVision(t)
SetVision(vision)
defer func() {
vision = nil
}()
ctx := context.Background()
t.Run("Image File with Vision Enabled", 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{}{
"vision": true,
"model": "gpt-4-vision-preview",
})
assert.NoError(t, err)
assert.NotNil(t, fileResp)
if fileResp.URL == "" && fileResp.Description == "" {
t.Error("Either URL or Description should be set when vision is enabled")
}
})
t.Run("Image File with Vision Disabled", 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{}{
"vision": false,
})
assert.NoError(t, err)
assert.NotNil(t, fileResp)
assert.Empty(t, fileResp.URL, "Vision URL should be empty when vision is disabled")
assert.Empty(t, fileResp.Description, "Vision Description should be empty when vision is disabled")
})
t.Run("Image File 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{}{
"vision": true,
"model": "gpt-4",
})
assert.NoError(t, err)
assert.NotNil(t, fileResp)
assert.Empty(t, fileResp.URL, "Vision URL should be empty for non-vision models")
assert.NotEmpty(t, fileResp.Description, "Vision Description should be set for non-vision models")
})
}
func TestDownload(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
ast := setupTestAssistant()
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")
})
}

229
neo/assistant/hooks.go Normal file
View file

@ -0,0 +1,229 @@
package assistant
import (
"context"
"fmt"
"time"
"github.com/gin-gonic/gin"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
)
// HookInit initialize the assistant
func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Init", context, input, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
response := &ResHookInit{}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["assistant_id"].(string); ok {
response.AssistantID = res
}
if res, ok := v["chat_id"].(string); ok {
response.ChatID = res
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
case string:
response.AssistantID = v
response.ChatID = context.ChatID
case nil:
response.AssistantID = ast.ID
response.ChatID = context.ChatID
}
return response, nil
}
// HookStream Handle streaming response from LLM
func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Stream", context, input, output, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
response := &ResHookStream{}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
// Custom silent from hook
if res, ok := v["silent"].(bool); ok {
response.Silent = res
}
case string:
response.Output = v
}
return response, nil
}
// HookDone Handle completion of assistant response
func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Done", context, input, output, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
response := &ResHookDone{
Input: input,
Output: output,
}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
case string:
response.Output = v
}
return response, nil
}
// HookFail Handle failure of assistant response
func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, output string, err error) (*ResHookFail, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, callErr := ast.call(ctx, "Fail", context, input, output, err.Error(), c.Writer)
if callErr != nil {
if callErr.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, callErr
}
response := &ResHookFail{
Input: input,
Output: output,
Error: err.Error(),
}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
}
if res, ok := v["error"].(string); ok {
response.Error = res
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
response.Next.Action = name
}
if payload, ok := res["payload"].(map[string]interface{}); ok {
response.Next.Payload = payload
}
}
case string:
response.Output = v
}
return response, nil
}
// createTimeoutContext creates a timeout context with 5 seconds timeout
func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) {
ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second)
return ctx, cancel
}
// Call the script method
func (ast *Assistant) call(ctx context.Context, method string, context chatctx.Context, args ...any) (interface{}, error) {
if ast.Script == nil {
return nil, nil
}
scriptCtx, err := ast.Script.NewContext(context.Sid, nil)
if err != nil {
return nil, err
}
defer scriptCtx.Close()
// Check if the method exists
if !scriptCtx.Global().Has(method) {
return nil, fmt.Errorf(HookErrorMethodNotFound)
}
// Create done channel for handling cancellation
done := make(chan struct{})
var result interface{}
var callErr error
go func() {
defer close(done)
// Call the method
args = append([]interface{}{context.Map()}, args...)
result, callErr = scriptCtx.Call(method, args...)
}()
// Wait for either context cancellation or method completion
select {
case <-ctx.Done():
scriptCtx.Close() // Force close the script context
return nil, ctx.Err()
case <-done:
return result, callErr
}
}

4
neo/assistant/load_test.go
View file

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

87
neo/assistant/types.go
View file

@ -5,17 +5,64 @@ import (
"io"
"mime/multipart"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/rag/driver"
v8 "github.com/yaoapp/gou/runtime/v8"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
api "github.com/yaoapp/yao/openai"
)
const (
// HookErrorMethodNotFound is the error message for method not found
HookErrorMethodNotFound = "method not found"
)
// API the assistant API interface
type API interface {
Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error
Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error
Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error)
Download(ctx context.Context, fileID string) (*FileResponse, error)
ReadBase64(ctx context.Context, fileID string) (string, error)
Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error
HookInit(c *gin.Context, ctx chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error)
}
// ResHookInit the response of the init hook
type ResHookInit struct {
AssistantID string `json:"assistant_id,omitempty"`
ChatID string `json:"chat_id,omitempty"`
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Options map[string]interface{} `json:"options,omitempty"`
}
// ResHookStream the response of the stream hook
type ResHookStream struct {
Silent bool `json:"silent,omitempty"` // Whether to suppress the output
Next *NextAction `json:"next,omitempty"` // The next action
Output string `json:"output,omitempty"` // The output
}
// ResHookDone the response of the done hook
type ResHookDone struct {
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Output string `json:"output,omitempty"`
}
// ResHookFail the response of the fail hook
type ResHookFail struct {
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Output string `json:"output,omitempty"`
Error string `json:"error,omitempty"`
}
// NextAction the next action
type NextAction struct {
Action string `json:"action"`
Payload map[string]interface{} `json:"payload,omitempty"`
}
// RAG the RAG interface
@ -70,13 +117,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

49
neo/context.go
View file

@ -1,49 +0,0 @@
package neo
import (
"context"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
)
// NewContext create a new context
func NewContext(sid, cid, payload string) Context {
ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid}
if payload == "" {
return ctx
}
err := jsoniter.Unmarshal([]byte(payload), &ctx)
if err != nil {
log.Error("%s", err.Error())
}
return ctx
}
// NewContextWithCancel create a new context with cancel
func NewContextWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
ctx := NewContext(sid, cid, payload)
return ContextWithCancel(ctx)
}
// NewContextWithTimeout create a new context with timeout
func NewContextWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
ctx := NewContext(sid, cid, payload)
return ContextWithTimeout(ctx, timeout)
}
// ContextWithCancel create a new context
func ContextWithCancel(parent Context) (Context, context.CancelFunc) {
new, cancel := context.WithCancel(parent.Context)
parent.Context = new
return parent, cancel
}
// ContextWithTimeout create a new context
func ContextWithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
new, cancel := context.WithTimeout(parent.Context, timeout)
parent.Context = new
return parent, cancel
}

122
neo/context/context.go Normal file
View file

@ -0,0 +1,122 @@
package context
import (
"context"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
)
// Context the context
type Context struct {
context.Context
Sid string `json:"sid" yaml:"-"` // Session ID
ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat
AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
Stack string `json:"stack,omitempty"`
Path string `json:"pathname,omitempty"`
FormData map[string]interface{} `json:"formdata,omitempty"`
Field *Field `json:"field,omitempty"`
Namespace string `json:"namespace,omitempty"`
Config map[string]interface{} `json:"config,omitempty"`
Signal interface{} `json:"signal,omitempty"`
Upload *FileUpload `json:"upload,omitempty"`
}
// Field the context field
type Field struct {
Name string `json:"name,omitempty"`
Type string `json:"type,omitempty"`
Bind string `json:"bind,omitempty"`
Props map[string]interface{} `json:"props,omitempty"`
Children []interface{} `json:"children,omitempty"`
}
// FileUpload the file upload
type FileUpload struct {
Name string `json:"name,omitempty"`
Type string `json:"type,omitempty"`
Size int64 `json:"size,omitempty"`
TempFile string `json:"temp_file,omitempty"`
}
// New create a new context
func New(sid, cid, payload string) Context {
ctx := Context{Context: context.Background(), Sid: sid, ChatID: cid}
if payload == "" {
return ctx
}
err := jsoniter.Unmarshal([]byte(payload), &ctx)
if err != nil {
log.Error("%s", err.Error())
}
return ctx
}
// NewWithCancel create a new context with cancel
func NewWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
ctx := New(sid, cid, payload)
return WithCancel(ctx)
}
// NewWithTimeout create a new context with timeout
func NewWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
ctx := New(sid, cid, payload)
return WithTimeout(ctx, timeout)
}
// WithCancel create a new context
func WithCancel(parent Context) (Context, context.CancelFunc) {
new, cancel := context.WithCancel(parent.Context)
parent.Context = new
return parent, cancel
}
// WithTimeout create a new context
func WithTimeout(parent Context, timeout time.Duration) (Context, context.CancelFunc) {
new, cancel := context.WithTimeout(parent.Context, timeout)
parent.Context = new
return parent, cancel
}
// Map the context to a map
func (ctx *Context) Map() map[string]interface{} {
data := map[string]interface{}{
"sid": ctx.Sid,
}
if ctx.ChatID != "" {
data["chat_id"] = ctx.ChatID
}
if ctx.AssistantID != "" {
data["assistant_id"] = ctx.AssistantID
}
if ctx.Stack != "" {
data["stack"] = ctx.Stack
}
if ctx.Path != "" {
data["pathname"] = ctx.Path
}
if len(ctx.FormData) > 0 {
data["formdata"] = ctx.FormData
}
if ctx.Field != nil {
data["field"] = ctx.Field
}
if ctx.Namespace != "" {
data["namespace"] = ctx.Namespace
}
if len(ctx.Config) > 0 {
data["config"] = ctx.Config
}
if ctx.Signal != nil {
data["signal"] = ctx.Signal
}
if ctx.Upload != nil {
data["upload"] = ctx.Upload
}
return data
}

7
neo/hooks.go
View file

@ -7,10 +7,11 @@ import (
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/process"
chatctx "github.com/yaoapp/yao/neo/context"
)
// HookCreate create the assistant
func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
// Default assistant
assistantID := neo.Use
@ -69,7 +70,7 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi
}
// HookPrepare executes the prepare hook before AI is called
func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
func (neo *DSL) HookPrepare(ctx chatctx.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
if neo.Prepare == "" {
return messages, nil
}
@ -114,7 +115,7 @@ func (neo *DSL) HookPrepare(ctx Context, messages []map[string]interface{}) ([]m
}
// HookWrite executes the write hook when response is received from AI
func (neo *DSL) HookWrite(ctx Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) {
func (neo *DSL) HookWrite(ctx chatctx.Context, messages []map[string]interface{}, response map[string]interface{}, content string, writer *gin.ResponseWriter) ([]map[string]interface{}, error) {
if neo.Write == "" {
return []map[string]interface{}{response}, nil
}

64
neo/message/message.go
View file

@ -22,6 +22,8 @@ type Message struct {
IsDone bool `json:"done,omitempty"`
Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend
Attachments []Attachment `json:"attachments,omitempty"` // File attachments
Role string `json:"role,omitempty"` // user, assistant, system ...
Name string `json:"name,omitempty"` // name for the message
Data map[string]interface{} `json:"-"`
}
@ -134,12 +136,74 @@ func (m *Message) Error(message interface{}) *Message {
return m
}
// SetContent set the content
func (m *Message) SetContent(content string) *Message {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
m.Text = err.Error() + "\n" + content
return m
}
*m = msg
} else {
m.Text = content
m.Type = "text"
}
return m
}
// Content get the content
func (m *Message) Content() string {
content := map[string]interface{}{"text": m.Text}
if m.Attachments != nil {
content["attachments"] = m.Attachments
}
if m.Type != "" {
content["type"] = m.Type
}
contentRaw, _ := jsoniter.MarshalToString(content)
return contentRaw
}
// ToMap convert to map
func (m *Message) ToMap() map[string]interface{} {
return map[string]interface{}{
"content": m.Content(),
"role": m.Role,
"name": m.Name,
}
}
// Map set from map
func (m *Message) Map(msg map[string]interface{}) *Message {
if msg == nil {
return m
}
// Content {"text": "xxxx", "attachments": ... }
if content, ok := msg["content"].(string); ok {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
m.Text = err.Error() + "\n" + content
return m
}
*m = msg
} else {
m.Text = content
m.Type = "text"
}
}
if role, ok := msg["role"].(string); ok {
m.Role = role
}
if name, ok := msg["name"].(string); ok {
m.Name = name
}
if text, ok := msg["text"].(string); ok {
m.Text = text
}

212
neo/neo.go
View file

@ -4,41 +4,25 @@ import (
"fmt"
"os"
"strings"
"sync"
"github.com/gin-gonic/gin"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/neo/assistant"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
)
// Lock the assistant list
var lock sync.Mutex = sync.Mutex{}
// Answer reply the message
func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error {
messages, err := neo.chatMessages(ctx, question)
if err != nil {
msg := message.New().Error(err).Done()
msg.Write(c.Writer)
return err
func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error {
var err error
var ast assistant.API = neo.Assistant
if ctx.AssistantID != "" {
ast, err = neo.Select(ctx.AssistantID)
if err != nil {
return err
}
}
// Get the assistant_id, chat_id
res, err := neo.HookCreate(ctx, messages, c)
if err != nil {
msg := message.New().Error(err).Done()
msg.Write(c.Writer)
return err
}
// Select Assistant
ast, err := neo.Select(res.AssistantID)
if err != nil {
return err
}
// Chat with AI
return neo.chat(ast, ctx, messages, c)
return ast.Execute(c, ctx, question, nil)
}
// Select select an assistant
@ -50,7 +34,7 @@ func (neo *DSL) Select(id string) (assistant.API, error) {
}
// GeneratePrompts generate prompts for the AI assistant
func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) {
func (neo *DSL) GeneratePrompts(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) {
prompts := `
Optimize the prompts for the AI assistant
1. Optimize prompts based on the user's input
@ -68,13 +52,14 @@ func (neo *DSL) GeneratePrompts(ctx Context, input string, c *gin.Context, silen
}
// GenerateChatTitle generate the chat title
func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, silent ...bool) (string, error) {
func (neo *DSL) GenerateChatTitle(ctx chatctx.Context, input string, c *gin.Context, silent ...bool) (string, error) {
prompts := `
Help me generate a title for the chat
1. The title should be a short and concise description of the chat.
2. The title should be a single sentence.
3. The title should be in same language as the chat.
4. The title should be no more than 50 characters.
5. ANSWER ONLY THE TITLE CONTENT, FOR EXAMPLE: Chat with AI is a valid title, but "Chat with AI" is not a valid title.
`
isSilent := false
if len(silent) > 0 {
@ -84,7 +69,7 @@ func (neo *DSL) GenerateChatTitle(ctx Context, input string, c *gin.Context, sil
}
// GenerateWithAI generate content with AI, type can be "title", "prompts", etc.
func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType string, systemPrompt string, c *gin.Context, silent bool) (string, error) {
messages := []map[string]interface{}{
{"role": "system", "content": systemPrompt},
{
@ -119,7 +104,11 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy
// Chat with AI in background
go func() {
err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int {
msgList := make([]message.Message, len(messages))
for i, msg := range messages {
msgList[i] = *message.New().Map(msg)
}
err := ast.Chat(c.Request.Context(), msgList, neo.Option, func(data []byte) int {
select {
case <-clientBreak:
return 0 // break
@ -187,7 +176,7 @@ func (neo *DSL) GenerateWithAI(ctx Context, input string, messageType string, sy
}
// Upload upload a file
func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
func (neo *DSL) Upload(ctx chatctx.Context, c *gin.Context) (*assistant.File, error) {
// Get the file
tmpfile, err := c.FormFile("file")
if err != nil {
@ -212,29 +201,30 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
}
// Get file info
ctx.Upload = &FileUpload{
Bytes: int(tmpfile.Size),
Name: tmpfile.Filename,
ContentType: tmpfile.Header.Get("Content-Type"),
Option: option,
ctx.Upload = &chatctx.FileUpload{
Name: tmpfile.Filename,
Type: tmpfile.Header.Get("Content-Type"),
Size: tmpfile.Size,
TempFile: tmpfile.Filename,
}
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)
}
// Download downloads a file
func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, error) {
func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileResponse, error) {
// Get file_id from query string
fileID := c.Query("file_id")
if fileID == "" {
@ -256,133 +246,3 @@ func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse,
// Download file using the assistant
return ast.Download(ctx.Context, fileID)
}
// chat chat with AI
func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error {
if ast == nil {
msg := message.New().Error("assistant is not initialized").Done()
msg.Write(c.Writer)
return fmt.Errorf("assistant is not initialized")
}
clientBreak := make(chan bool, 1)
done := make(chan bool, 1)
content := []byte{}
// Chat with AI in background
go func() {
err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int {
select {
case <-clientBreak:
return 0 // break
default:
msg := message.NewOpenAI(data)
if msg == nil {
return 1 // continue
}
// Handle error
if msg.Type == "error" {
value := msg.String()
message.New().Error(value).Done().Write(c.Writer)
return 0 // break
}
// Append content and send message
content = msg.Append(content)
value := msg.String()
if value != "" {
message.New().
Map(map[string]interface{}{
"text": value,
"done": msg.IsDone,
}).
Write(c.Writer)
}
// Complete the stream
if msg.IsDone {
if value == "" {
msg.Write(c.Writer)
}
done <- true
return 0 // break
}
return 1 // continue
}
})
if err != nil {
log.Error("Chat error: %s", err.Error())
message.New().Error(err).Done().Write(c.Writer)
}
// Save chat history
if len(content) > 0 {
neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages)
}
done <- true
}()
// Wait for completion or client disconnect
select {
case <-done:
return nil
case <-c.Writer.CloseNotify():
clientBreak <- true
return nil
}
}
// chatMessages get the chat messages
func (neo *DSL) chatMessages(ctx Context, content ...string) ([]map[string]interface{}, error) {
history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID)
if err != nil {
return nil, err
}
messages := []map[string]interface{}{}
messages = append(messages, history...)
if len(content) == 0 {
return messages, nil
}
// Add user message
messages = append(messages, map[string]interface{}{"role": "user", "content": content[0], "name": ctx.Sid})
return messages, nil
}
// saveHistory save the history
func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages []map[string]interface{}) {
if len(content) > 0 && sid != "" && len(messages) > 0 {
err := neo.Store.SaveHistory(
sid,
[]map[string]interface{}{
{"role": "user", "content": messages[len(messages)-1]["content"], "name": sid},
{"role": "assistant", "content": string(content), "name": sid},
},
chatID,
nil,
)
if err != nil {
log.Error("Save history error: %s", err.Error())
}
}
}
// 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 {
if !message.New().Map(msg).Write(w) {
return fmt.Errorf("failed to write message to stream")
}
return nil
}
return fmt.Errorf("invalid message data type")
}

18
neo/types.go
View file

@ -1,8 +1,6 @@
package neo
import (
"context"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/rag"
@ -48,22 +46,6 @@ type Mention struct {
Type string `json:"type,omitempty"`
}
// Context the context
type Context struct {
Sid string `json:"sid" yaml:"-"` // Session ID
ChatID string `json:"chat_id,omitempty"` // Chat ID, use to select chat
AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, use to select assistant
Stack string `json:"stack,omitempty"`
Path string `json:"pathname,omitempty"`
FormData map[string]interface{} `json:"formdata,omitempty"`
Field *Field `json:"field,omitempty"`
Namespace string `json:"namespace,omitempty"`
Config map[string]interface{} `json:"config,omitempty"`
Signal interface{} `json:"signal,omitempty"`
Upload *FileUpload `json:"upload,omitempty"`
context.Context `json:"-" yaml:"-"`
}
// Field the context field
type Field struct {
Name string `json:"name,omitempty"`

10
neo/vision/driver/openai/model.go
View file

@ -52,11 +52,17 @@ func New(options map[string]interface{}) (*Model, error) {
}
// Analyze analyze image using OpenAI vision model
func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error) {
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")
}
// Use default prompt if none provided
userPrompt := model.Prompt
if len(prompt) > 0 && prompt[0] != "" {
userPrompt = prompt[0]
}
// Check if fileID is a URL or base64 data
var imageURL string
if strings.HasPrefix(fileID, "data:image/") {
@ -103,7 +109,7 @@ func (model *Model) Analyze(ctx context.Context, fileID string, prompt string) (
"content": []map[string]interface{}{
{
"type": "text",
"text": prompt,
"text": userPrompt,
},
{
"type": "image_url",

45
neo/vision/driver/openai/model_test.go
View file

@ -146,4 +146,49 @@ func TestOpenAIModel(t *testing.T) {
assert.Error(t, err)
assert.Contains(t, err.Error(), "OpenAI API error")
})
t.Run("Analyze with Default Prompt", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": "Default test prompt",
})
assert.NoError(t, err)
// Use base64 image data without providing a prompt
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with Custom Prompt Overriding Default", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": "Default test prompt",
})
assert.NoError(t, err)
// Use base64 image data with custom prompt
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
t.Run("Analyze with Empty Custom Prompt", func(t *testing.T) {
model, err := New(map[string]interface{}{
"api_key": os.Getenv("OPENAI_API_KEY"),
"model": os.Getenv("VISION_MODEL"),
"prompt": "Default test prompt",
})
assert.NoError(t, err)
// Use base64 image data with empty prompt (should use default)
result, err := model.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result["description"])
})
}

4
neo/vision/driver/types.go
View file

@ -32,7 +32,9 @@ type Storage interface {
// Model the vision model interface
type Model interface {
Analyze(ctx context.Context, fileID string, prompt string) (map[string]interface{}, error)
// Analyze analyzes an image file
// If prompt is empty, it will use the default prompt from model.options.prompt
Analyze(ctx context.Context, fileID string, prompt ...string) (map[string]interface{}, error)
}
// Response the vision response

43
neo/vision/vision.go
View file

@ -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)
}
@ -75,7 +104,7 @@ func (v *Vision) Upload(ctx context.Context, filename string, reader io.Reader,
}
// Analyze analyze image using vision model
func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*driver.Response, error) {
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")
}
@ -92,7 +121,7 @@ func (v *Vision) Analyze(ctx context.Context, fileID string, prompt string) (*dr
}
}
result, err := v.model.Analyze(ctx, url, prompt)
result, err := v.model.Analyze(ctx, url, prompt...)
if err != nil {
return nil, err
}

90
neo/vision/vision_test.go
View file

@ -331,6 +331,96 @@ func TestVision(t *testing.T) {
assert.LessOrEqual(t, bounds.Dx(), MaxImageSize)
assert.LessOrEqual(t, bounds.Dy(), MaxImageSize)
})
t.Run("Analyze Image with Default Prompt", func(t *testing.T) {
// Create vision service with default prompt
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"),
"prompt": "Default test prompt",
},
},
}
vision, err := New(cfg)
assert.NoError(t, err)
// Use base64 data without providing a prompt
result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Description)
})
t.Run("Analyze Image with Custom Prompt", func(t *testing.T) {
// Create vision service with default prompt
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"),
"prompt": "Default test prompt",
},
},
}
vision, err := New(cfg)
assert.NoError(t, err)
// Use base64 data with custom prompt
result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "Custom test prompt")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Description)
})
t.Run("Analyze Image with Empty Custom Prompt", func(t *testing.T) {
// Create vision service with default prompt
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"),
"prompt": "Default test prompt",
},
},
}
vision, err := New(cfg)
assert.NoError(t, err)
// Use base64 data with empty prompt (should use default)
result, err := vision.Analyze(context.Background(), "data:image/png;base64,"+testImageBase64, "")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotEmpty(t, result.Description)
})
}
func createTestVision(baseURL string) (*Vision, error) {