Remove hooks.go file and refactor Load function in load.go to streamline assistant initialization

- Deleted hooks.go as it was no longer needed.
- Updated Load function in load.go to set default assistant name to "mohe" and adjusted store settings.
- Removed unused RAG and Vision initialization functions to simplify the codebase.
- Refactored defaultAssistant function to improve error handling for missing default assistant.
This commit is contained in:
Max 2025-05-24 18:37:51 +08:00
parent ea053856c9
commit c087ed1760
4 changed files with 111 additions and 300 deletions

View file

@ -1,165 +0,0 @@
package neo
import (
"context"
"time"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/process"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
)
// HookCreate create the assistant
func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) {
// Default assistant
assistantID := neo.Use.Default
if ctx.AssistantID != "" {
assistantID = ctx.AssistantID
}
// Empty hook
if neo.Create == "" {
return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil
}
// Create a context with 10 second timeout
timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
defer cancel()
p, err := process.Of(neo.Create, ctx, messages, c.Writer)
if err != nil {
return CreateResponse{}, err
}
err = p.WithContext(timeoutCtx).Execute()
if err != nil {
return CreateResponse{}, err
}
defer p.Release()
// Check if context was canceled
if timeoutCtx.Err() != nil {
return CreateResponse{}, timeoutCtx.Err()
}
value := p.Value()
switch v := value.(type) {
case CreateResponse:
return v, nil
case map[string]interface{}:
if id, ok := v["assistant_id"].(string); ok {
assistantID = id
}
chatID := ""
if id, ok := v["chat_id"].(string); ok {
chatID = id
}
if chatID == "" {
chatID = ctx.ChatID
}
// Messages fixed input
input := []message.Message{}
if vv, has := v["input"]; has {
bytes, err := jsoniter.Marshal(vv)
if err != nil {
return CreateResponse{}, err
}
err = jsoniter.Unmarshal(bytes, &input)
if err != nil {
return CreateResponse{}, err
}
}
return CreateResponse{AssistantID: assistantID, ChatID: chatID, Input: input}, nil
}
return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID, Input: nil}, nil
}
// HookPrepare executes the prepare hook before AI is called
func (neo *DSL) HookPrepare(ctx chatctx.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) {
if neo.Prepare == "" {
return messages, nil
}
// Create a context with 10 second timeout
timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
p, err := process.Of(neo.Prepare, ctx, messages)
if err != nil {
return nil, err
}
err = p.WithContext(timeoutCtx).Execute()
if err != nil {
return nil, err
}
defer p.Release()
// Check if context was canceled
if timeoutCtx.Err() != nil {
return nil, timeoutCtx.Err()
}
value := p.Value()
if value == nil {
return messages, nil
}
var result []map[string]interface{}
bytes, err := jsoniter.Marshal(value)
if err != nil {
return nil, err
}
err = jsoniter.Unmarshal(bytes, &result)
if err != nil {
return nil, err
}
return result, nil
}
// HookWrite executes the write hook when response is received from AI
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
}
p, err := process.Of(neo.Write, ctx, messages, response, content, writer)
if err != nil {
return nil, err
}
err = p.WithContext(ctx).Execute()
if err != nil {
return nil, err
}
defer p.Release()
value := p.Value()
if value == nil {
return []map[string]interface{}{response}, nil
}
var result []map[string]interface{}
bytes, err := jsoniter.Marshal(value)
if err != nil {
return nil, err
}
err = jsoniter.Unmarshal(bytes, &result)
if err != nil {
return nil, err
}
return result, nil
}

View file

@ -4,16 +4,11 @@ import (
"fmt"
"path/filepath"
"github.com/fatih/color"
"github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/rag"
"github.com/yaoapp/yao/neo/store"
"github.com/yaoapp/yao/neo/vision"
"github.com/yaoapp/yao/neo/vision/driver"
)
// Neo the neo AI assistant
@ -23,10 +18,8 @@ var Neo *DSL
func Load(cfg config.Config) error {
setting := DSL{
ID: "neo",
Prompts: []assistant.Prompt{},
Option: map[string]interface{}{},
Allows: []string{},
ID: "neo",
Allows: []string{},
StoreSetting: store.Setting{
Prefix: "yao_neo_",
Connector: "default",
@ -44,12 +37,12 @@ func Load(cfg config.Config) error {
}
if setting.StoreSetting.MaxSize == 0 {
setting.StoreSetting.MaxSize = 100
setting.StoreSetting.MaxSize = 20 // default is 20
}
// Default Assistant
// Default Assistant, Neo is the developer name, Mohe is the brand name of the assistant
if setting.Use == nil {
setting.Use = &Use{Default: "neo"}
setting.Use = &Use{Default: "mohe"} // Neo is the developer name, Mohe is the brand name of the assistant
}
// Title Assistant
@ -70,12 +63,6 @@ func Load(cfg config.Config) error {
return err
}
// Initialize RAG
initRAG()
// Initialize Vision
initVision()
// Initialize Assistant
err = initAssistant()
if err != nil {
@ -85,21 +72,6 @@ func Load(cfg config.Config) error {
return nil
}
// initRAG initialize the RAG instance
func initRAG() {
if Neo.RAGSetting.Engine.Driver == "" {
return
}
instance, err := rag.New(Neo.RAGSetting)
if err != nil {
color.Red("[Neo] Failed to initialize RAG: %v", err)
log.Error("[Neo] Failed to initialize RAG: %v", err)
return
}
Neo.RAG = instance
}
// initStore initialize the store
func initStore() error {
@ -131,45 +103,12 @@ func initStore() error {
return fmt.Errorf("%s store connector %s not support", Neo.ID, Neo.StoreSetting.Connector)
}
// initVision initialize the Vision instance
func initVision() {
if Neo.VisionSetting.Storage.Driver == "" {
return
}
cfg := &driver.Config{
Storage: Neo.VisionSetting.Storage,
Model: Neo.VisionSetting.Model,
}
instance, err := vision.New(cfg)
if err != nil {
color.Red("[Neo] Failed to initialize Vision: %v", err)
log.Error("[Neo] Failed to initialize Vision: %v", err)
return
}
Neo.Vision = instance
}
// initAssistant initialize the assistant
func initAssistant() error {
// Set Storage
assistant.SetStorage(Neo.Store)
// Assistant RAG
if Neo.RAG != nil {
assistant.SetRAG(
Neo.RAG.Engine(),
Neo.RAG.FileUpload(),
Neo.RAG.Vectorizer(),
assistant.RAGSetting{
IndexPrefix: Neo.RAGSetting.IndexPrefix,
},
)
}
// Assistant Vision
if Neo.Vision != nil {
assistant.SetVision(Neo.Vision)
@ -179,9 +118,6 @@ func initAssistant() error {
assistant.SetConnectorSettings(Neo.Connectors)
}
// Default Connector
assistant.SetConnector(Neo.Connector)
// Load Built-in Assistants
err := assistant.LoadBuiltIn()
if err != nil {
@ -200,14 +136,8 @@ func initAssistant() error {
// defaultAssistant get the default assistant
func defaultAssistant() (*assistant.Assistant, error) {
if Neo.Use != nil && Neo.Use.Default != "" {
return assistant.Get(Neo.Use.Default)
if Neo.Use == nil || Neo.Use.Default == "" {
return nil, fmt.Errorf("default assistant not found")
}
name := Neo.Name
if name == "" {
name = "Neo"
}
return assistant.GetByConnector(Neo.Connector, name)
return assistant.Get(Neo.Use.Default)
}

View file

@ -89,13 +89,13 @@ func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileRe
}
// Get assistant_id from context or query
res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c)
if err != nil {
return nil, err
}
// res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c)
// if err != nil {
// return nil, err
// }
// Select Assistant
ast, err := neo.Select(res.AssistantID)
ast, err := neo.Select(neo.Use.Default)
if err != nil {
return nil, err
}

View file

@ -3,75 +3,121 @@ package neo
import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/message"
"github.com/yaoapp/yao/neo/rag"
"github.com/yaoapp/yao/neo/store"
"github.com/yaoapp/yao/neo/vision"
"github.com/yaoapp/yao/neo/vision/driver"
)
// DSL AI assistant
type DSL struct {
ID string `json:"-" yaml:"-"`
Name string `json:"name,omitempty" yaml:"name,omitempty"`
Use *Use `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default, title, prompt
Guard string `json:"guard,omitempty" yaml:"guard,omitempty"`
Connector string `json:"connector" yaml:"connector"`
StoreSetting store.Setting `json:"store" yaml:"store"`
RAGSetting rag.Setting `json:"rag" yaml:"rag"`
VisionSetting VisionSetting `json:"vision" yaml:"vision"`
Option map[string]interface{} `json:"option" yaml:"option"`
Prepare string `json:"prepare,omitempty" yaml:"prepare,omitempty"`
Create string `json:"create,omitempty" yaml:"create,omitempty"`
Write string `json:"write,omitempty" yaml:"write,omitempty"`
Prompts []assistant.Prompt `json:"prompts,omitempty" yaml:"prompts,omitempty"`
Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"`
Connectors map[string]assistant.ConnectorSetting `json:"connectors,omitempty" yaml:"connectors,omitempty"`
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
Store store.Store `json:"-" yaml:"-"`
RAG *rag.RAG `json:"-" yaml:"-"`
Vision *vision.Vision `json:"-" yaml:"-"`
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
// Neo Global Settings
// ===============================
Use *Use `json:"use,omitempty" yaml:"use,omitempty"` // Which assistant to use default, title, prompt
StoreSetting store.Setting `json:"store" yaml:"store"` // The store setting of the assistant
AuthSetting *Auth `json:"auth,omitempty" yaml:"auth,omitempty"` // Authenticate Settings
UploadSetting *Upload `json:"upload,omitempty" yaml:"upload,omitempty"` // Upload Settings
KnowledgeSetting *Knowledge `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base Settings
// Global External Settings - connectors, tools, etc.
// ===============================
Connectors map[string]assistant.ConnectorSetting `json:"connectors,omitempty" yaml:"connectors,omitempty"` // The connectors of the assistant
// Neo API Settings
// ===============================s
Guard string `json:"guard,omitempty" yaml:"guard,omitempty"` // The guard of the assistant
Allows []string `json:"allows,omitempty" yaml:"allows,omitempty"` // The allowed domains of the assistant
// Internal
// ===============================
ID string `json:"-" yaml:"-"` // The id of the instance
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
Store store.Store `json:"-" yaml:"-"` // The store of the assistant
RAG *rag.RAG `json:"-" yaml:"-"`
Vision *vision.Vision `json:"-" yaml:"-"`
GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`
}
// Use the use setting for the assistant
// Use the default assistant settings
// ===============================
type Use struct {
Default string `json:"default,omitempty" yaml:"default,omitempty"`
Title string `json:"title,omitempty" yaml:"title,omitempty"`
Prompt string `json:"prompt,omitempty" yaml:"prompt,omitempty"`
Default string `json:"default,omitempty" yaml:"default,omitempty"` // The default assistant to use
Title string `json:"title,omitempty" yaml:"title,omitempty"` // The assistant for generating the topic title.
Prompt string `json:"prompt,omitempty" yaml:"prompt,omitempty"` // The assistant for generating the prompt.
Vision string `json:"vision,omitempty" yaml:"vision,omitempty"` // The assistant for generating the image/video description, if the assistant enable the vision and model not support vision, use the vision model to describe the image/video, and return the messages with the image/video's description.
Search string `json:"search,omitempty" yaml:"search,omitempty"` // The assistant for searching the knowledge, global web search. If not set, and the assistant enable the knowledge, it will search the result from the knowledge automatically.
Fetch string `json:"fetch,omitempty" yaml:"fetch,omitempty"` // The assistant for fetching the http/https/ftp/sftp/etc. file, and return the file's content. if not set, use the http process to fetch the file.
}
// VisionSetting the vision setting
type VisionSetting struct {
Storage driver.StorageConfig `json:"storage" yaml:"storage"`
Model driver.ModelConfig `json:"model" yaml:"model"`
// Auth Authenticate Settings
// ===============================
type Auth struct {
Models *AuthModels `json:"models,omitempty" yaml:"models,omitempty"` // The models of the user, it is used to handle the user, and the user is a user in the database. (Guest and User model must have the id and permission fields)
Fields *AuthFields `json:"fields,omitempty" yaml:"fields,omitempty"` // The fields of the user model, it is used to handle the user, and the user is a user in the database. (Guest and User model must have the id and permission fields)
SessionFields *AuthSessionFields `json:"session_fields,omitempty" yaml:"session_fields,omitempty"` // The session fields of the user, it is used to handle the user, and the user is a user in the database.
}
// Mention list
// AuthModels the auth model
type AuthModels struct {
User string `json:"user,omitempty" yaml:"user,omitempty"` // default is admin.user, The user model is a special model, it is used to handle the user, and the user is a user in the database.
Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // The guest model is a special model, it is used to handle the guest user, and the guest user is not a user in the database.
}
// AuthSessionFields the auth session field
type AuthSessionFields struct {
ID string `json:"id,omitempty" yaml:"id,omitempty"` // the field name of the user id, default is user_id
Roles string `json:"roles,omitempty" yaml:"roles,omitempty"` // the field name of the user roles, default is user_roles. the value must be an JSON array string.
Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // the field name of the guest user, default is guest
}
// AuthFields the auth field
type AuthFields struct {
ID string `json:"id,omitempty" yaml:"id,omitempty"` // the field name of the user id, default is id
Roles string `json:"roles,omitempty" yaml:"roles,omitempty"` // the field name of the user roles, default is roles, it must be an JSON field.
Permission string `json:"permission,omitempty" yaml:"permission,omitempty"` // the field name of the user permission, default is permission
}
// Upload the upload setting
// ===============================
type Upload struct {
Driver string `json:"driver" yaml:"driver"` // local, s3, default is local
Options map[string]interface{} `json:"options" yaml:"options"` // the options of the upload, it is used to configure the upload driver.
Compression bool `json:"compression,omitempty" yaml:"compression,omitempty"` // Compress the image/video to a smaller size, if the image/video is too large, it will be compressed to a smaller size.
ChunkSize string `json:"chunk_size,omitempty" yaml:"chunk_size,omitempty"` // the chunk size of the file, if the file is too large, it will be chunked into smaller chunks.
AllowedTypes []string `json:"allowed_types,omitempty" yaml:"allowed_types,omitempty"` // the allowed types of the file, if the file is not in the allowed types, it will be rejected.
}
// Knowledge base Settings
// ===============================
type Knowledge struct {
Vector KnowledgeVector `json:"vector" yaml:"vector"` // The vector database driver
Graph KnowledgeGraph `json:"graph" yaml:"graph"` // The graph database driver
Vectorizer KnowledgeVectorizer `json:"vectorizer" yaml:"vectorizer"` // The vectorizer driver
}
// KnowledgeVectorizer the knowledge vectorizer
type KnowledgeVectorizer struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// KnowledgeVector the knowledge vector
type KnowledgeVector struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// KnowledgeGraph the knowledge graph
type KnowledgeGraph struct {
Driver string `json:"driver" yaml:"driver"`
Options map[string]interface{} `json:"options" yaml:"options"`
}
// Mention Structure
// ===============================
type Mention struct {
ID string `json:"id"`
Name string `json:"name"`
Avatar string `json:"avatar,omitempty"`
Type string `json:"type,omitempty"`
}
// Field the context field
type Field struct {
Name string `json:"name,omitempty"`
Bind string `json:"bind,omitempty"`
}
// FileUpload the file upload info
type FileUpload struct {
Bytes int `json:"bytes,omitempty"` // If upload file, the file bytes
Name string `json:"name,omitempty"` // If upload
ContentType string `json:"content_type,omitempty"` // If upload file, the file content type
Option map[string]interface{} `json:"option,omitempty"` // If upload file, the upload option
}
// CreateResponse the response of the create hook
type CreateResponse struct {
AssistantID string `json:"assistant_id,omitempty"`
ChatID string `json:"chat_id,omitempty"`
Input []message.Message `json:"messages,omitempty"`
}