Refactor Neo API assistant to enhance vision support and message handling

- Updated the Assistant struct to include a vision capability flag and an init hook indicator, allowing for better management of vision-enabled functionalities.
- Refactored message handling throughout the assistant methods to utilize the chatMessage package, improving consistency and type safety.
- Enhanced the handleVision method to support dynamic vision processing options, including improved handling of image descriptions and uploads.
- Streamlined the requestMessages and withAttachments methods to better accommodate vision capabilities, ensuring proper integration with image handling.
- Improved the initialize method to check for vision support based on the model, enhancing the assistant's adaptability.

These changes improve the robustness and maintainability of the Neo API, paving the way for enhanced assistant functionalities and better integration of vision capabilities.
This commit is contained in:
Max 2025-01-17 15:16:26 +08:00
parent 5e5524293f
commit 17c18e7d56
6 changed files with 153 additions and 83 deletions

View file

@ -7,10 +7,11 @@ import (
"strings"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/utils"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
chatMessage "github.com/yaoapp/yao/neo/message"
)
@ -160,10 +161,10 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
}
// 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 {
func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}) error {
clientBreak := make(chan bool, 1)
done := make(chan bool, 1)
content := message.NewContent("text")
content := chatMessage.NewContent("text")
// Chat with AI in background
go func() {
@ -190,11 +191,11 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
func (ast *Assistant) streamChat(
c *gin.Context,
ctx chatctx.Context,
messages []message.Message,
messages []chatMessage.Message,
options map[string]interface{},
clientBreak chan bool,
done chan bool,
content *message.Content) error {
content *chatMessage.Content) error {
return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
select {
@ -276,7 +277,7 @@ func (ast *Assistant) streamChat(
// }
// Call HookDone
content.SetStatus(message.ContentStatusDone)
content.SetStatus(chatMessage.ContentStatusDone)
res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function")
if hookErr == nil && res != nil {
if res.Output != "" {
@ -316,7 +317,7 @@ func (ast *Assistant) streamChat(
}
// saveChatHistory saves the chat history if storage is available
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content *message.Content) {
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, content *chatMessage.Content) {
if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 {
storage.SaveHistory(
ctx.Sid,
@ -352,21 +353,21 @@ func (ast *Assistant) withOptions(options map[string]interface{}) map[string]int
return options
}
func (ast *Assistant) withPrompts(messages []message.Message) []message.Message {
func (ast *Assistant) withPrompts(messages []chatMessage.Message) []chatMessage.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}))
messages = append(messages, *chatMessage.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{}
func (ast *Assistant) withHistory(ctx chatctx.Context, input string) ([]chatMessage.Message, error) {
messages := []chatMessage.Message{}
messages = ast.withPrompts(messages)
if storage != nil {
history, err := storage.GetHistory(ctx.Sid, ctx.ChatID)
@ -376,17 +377,17 @@ func (ast *Assistant) withHistory(ctx chatctx.Context, input string) ([]message.
// Add history messages
for _, h := range history {
messages = append(messages, *message.New().Map(h))
messages = append(messages, *chatMessage.New().Map(h))
}
}
// Add user message
messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": input, "name": ctx.Sid}))
messages = append(messages, *chatMessage.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 []message.Message, option map[string]interface{}, cb func(data []byte) int) error {
func (ast *Assistant) Chat(ctx context.Context, messages []chatMessage.Message, option map[string]interface{}, cb func(data []byte) int) error {
if ast.openai == nil {
return fmt.Errorf("openai is not initialized")
}
@ -404,27 +405,10 @@ func (ast *Assistant) Chat(ctx context.Context, messages []message.Message, opti
return nil
}
func (ast *Assistant) requestMessages(ctx context.Context, messages []message.Message) ([]map[string]interface{}, error) {
func (ast *Assistant) requestMessages(ctx context.Context, messages []chatMessage.Message) ([]map[string]interface{}, error) {
newMessages := []map[string]interface{}{}
// With Prompts
if ast.Prompts != nil {
for _, prompt := range ast.Prompts {
msg := map[string]interface{}{
"role": prompt.Role,
"content": prompt.Content,
}
name := ast.Name
if prompt.Name != "" {
name = prompt.Name
}
msg["name"] = name
newMessages = append(newMessages, msg)
}
}
length := len(messages)
for index, message := range messages {
role := message.Role
if role == "" {
@ -454,12 +438,24 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []message.Me
}
newMessage["content"] = msg.Text
if msg.Attachments != nil {
content, err := ast.withAttachments(ctx, msg)
if message.Attachments != nil {
contents, err := ast.withAttachments(ctx, &message)
if err != nil {
return nil, fmt.Errorf("with attachments error: %s", err.Error())
}
newMessage["content"] = content
// if current assistant is vision capable, add the contents directly
if ast.vision {
newMessage["content"] = contents
continue
}
// If current assistant is not vision capable, add the description of the image
if contents != nil {
for _, content := range contents {
newMessages = append(newMessages, content)
}
}
}
}
@ -470,10 +466,27 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []message.Me
func (ast *Assistant) withAttachments(ctx context.Context, msg *chatMessage.Message) ([]map[string]interface{}, error) {
contents := []map[string]interface{}{{"type": "text", "text": msg.Text}}
if !ast.vision {
contents = []map[string]interface{}{{"role": "user", "content": msg.Text}}
}
images := []string{}
for _, attachment := range msg.Attachments {
if strings.HasPrefix(attachment.ContentType, "image/") {
images = append(images, attachment.FileID)
if ast.vision {
images = append(images, attachment.URL)
continue
}
// If the current assistant is not vision capable, add the description of the image
raw, err := jsoniter.MarshalToString(attachment)
if err != nil {
return nil, fmt.Errorf("marshal attachment error: %s", err.Error())
}
contents = append(contents, map[string]interface{}{
"role": "system",
"content": raw,
})
}
}
@ -481,20 +494,40 @@ func (ast *Assistant) withAttachments(ctx context.Context, msg *chatMessage.Mess
return contents, nil
}
for _, image := range images {
bytes64, err := ast.ReadBase64(ctx, image)
if err != nil {
return nil, fmt.Errorf("read base64 error: %s", err.Error())
// If the current assistant is vision capable, add the image to the contents directly
if ast.vision {
for _, url := range images {
// If the image is already a URL, add it directly
if strings.HasPrefix(url, "http") {
contents = append(contents, map[string]interface{}{
"type": "image_url",
"image_url": map[string]string{
"url": url,
},
})
continue
}
// Read base64
bytes64, err := ast.ReadBase64(ctx, url)
if err != nil {
return nil, fmt.Errorf("read base64 error: %s", err.Error())
}
contents = append(contents, map[string]interface{}{
"type": "image_url",
"image_url": map[string]string{
"url": fmt.Sprintf("data:image/jpeg;base64,%s", bytes64),
},
})
}
contents = append(contents, map[string]interface{}{
"type": "image_url",
"image_url": map[string]string{
"url": fmt.Sprintf("data:image/jpeg;base64,%s", bytes64),
},
})
utils.Dump(contents)
return contents, nil
}
// If the current assistant is not vision capable, add the description of the image
return contents, nil
}

View file

@ -197,12 +197,22 @@ func (ast *Assistant) handleRAG(ctx context.Context, file *File, reader io.Reade
// 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 {
handleVision := false
if vv, has := option["vision"]; has {
switch v := vv.(type) {
case bool:
handleVision = v
case string:
handleVision = v == "true" || v == "1" || v == "yes" || v == "on" || v == "enable"
}
}
if !handleVision {
return nil
}
@ -211,12 +221,6 @@ func (ast *Assistant) handleVision(ctx context.Context, file *File, option map[s
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 {
@ -237,43 +241,48 @@ func (ast *Assistant) handleVision(ctx context.Context, file *File, option map[s
return fmt.Errorf("read file error: %s", err.Error())
}
if VisionCapableModels[model] {
// The model is vision capable
if ast.vision {
// 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
return nil
}
// 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["description"].(string); ok {
file.Description = desc
} else if desc, ok := result.Description["text"].(string); ok {
file.Description = desc
} 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)
}
// Convert the entire description to JSON string as fallback
bytes, err := jsoniter.Marshal(result.Description)
if err == nil {
file.Description = string(bytes)
}
}
return nil
}

View file

@ -517,5 +517,25 @@ func (ast *Assistant) initialize() error {
return err
}
ast.openai = api
// Check if the assistant supports vision
model := api.Model()
if v, ok := ast.Options["model"].(string); ok {
model = strings.TrimLeft(v, "moapi:")
}
if _, ok := VisionCapableModels[model]; ok {
ast.vision = true
}
// Check if the assistant has an init hook
if ast.Script != nil {
scriptCtx, err := ast.Script.NewContext("", nil)
if err != nil {
return err
}
defer scriptCtx.Close()
ast.initHook = scriptCtx.Global().Has("init")
}
return nil
}

View file

@ -126,6 +126,8 @@ type Assistant struct {
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
openai *api.OpenAI // OpenAI API
vision bool // Whether this assistant supports vision
initHook bool // Whether this assistant has an init hook
}
// VisionCapableModels list of LLM models that support vision capabilities

View file

@ -31,6 +31,7 @@ type Message struct {
type Attachment struct {
Name string `json:"name,omitempty"`
URL string `json:"url,omitempty"`
Description string `json:"description,omitempty"`
Type string `json:"type,omitempty"`
ContentType string `json:"content_type,omitempty"`
Bytes int64 `json:"bytes,omitempty"`

View file

@ -127,6 +127,11 @@ func NewMoapi(model string) (*OpenAI, error) {
}, nil
}
// Model get the model
func (openai OpenAI) Model() string {
return openai.model
}
// Completions Creates a completion for the provided prompt and parameters.
// https://platform.openai.com/docs/api-reference/completions/create
func (openai OpenAI) Completions(prompt interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) {