Enhance RAG configuration and assistant loading in Neo API
- Updated the SetRAG function to include RAG settings, allowing for more detailed configuration during assistant initialization. - Refactored the LoadBuiltIn function to streamline the saving process of assistants, improving error handling and code clarity. - Modified loadPrompts and loadScript functions to return timestamps, enabling tracking of creation and update times for assistants. - Enhanced the Assistant struct with created_at and updated_at fields, improving data management and traceability. - Overall improvements to the assistant loading process, ensuring better organization and maintainability of the codebase.
This commit is contained in:
Max
parent
973fb5dd3d
commit
af30bec4f0
5 changed files with 374 additions and 165 deletions
253
neo/assistant/assistant.go
Normal file
253
neo/assistant/assistant.go
Normal file
|
|
@ -0,0 +1,253 @@
|
|||
package assistant
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/fatih/color"
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
"github.com/yaoapp/gou/rag/driver"
|
||||
"github.com/yaoapp/kun/log"
|
||||
)
|
||||
|
||||
// Save save the assistant
|
||||
func (ast *Assistant) Save() error {
|
||||
if storage == nil {
|
||||
return fmt.Errorf("storage is not set")
|
||||
}
|
||||
|
||||
_, err := storage.SaveAssistant(ast.Map())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Update Index in background
|
||||
go func() {
|
||||
err := ast.UpdateIndex()
|
||||
if err != nil {
|
||||
log.Error("failed to update index for assistant %s: %s", ast.ID, err)
|
||||
color.Red("failed to update index for assistant %s: %s", ast.ID, err)
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateIndex update the index for RAG
|
||||
func (ast *Assistant) UpdateIndex() error {
|
||||
|
||||
// RAG is not enabled
|
||||
if rag == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if rag.Engine == nil {
|
||||
return fmt.Errorf("engine is not set")
|
||||
}
|
||||
|
||||
// Update Index
|
||||
index := fmt.Sprintf("%sassistants", rag.Setting.IndexPrefix)
|
||||
id := fmt.Sprintf("assistant_%s", ast.ID)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
|
||||
// Check if the index exists
|
||||
exists, err := rag.Engine.HasIndex(ctx, index)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Create the index if it does not exist
|
||||
if !exists {
|
||||
ctxCreate, cancelCreate := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancelCreate()
|
||||
err = rag.Engine.CreateIndex(ctxCreate, driver.IndexConfig{Name: index})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Check if the document exists
|
||||
exists, err = rag.Engine.HasDocument(ctx, index, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Check if the document is updated
|
||||
if exists {
|
||||
metadata, err := rag.Engine.GetMetadata(ctx, index, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if v, ok := metadata["updated_at"].(string); ok {
|
||||
updatedAt, err := stringToTimestamp(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if updatedAt >= ast.UpdatedAt {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update the index
|
||||
content, err := jsoniter.MarshalToString(ast.Map())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
metadata := map[string]interface{}{
|
||||
"assistant_id": ast.ID,
|
||||
"type": ast.Type,
|
||||
"name": ast.Name,
|
||||
"updated_at": fmt.Sprintf("%d", ast.UpdatedAt),
|
||||
}
|
||||
|
||||
return rag.Engine.IndexDoc(ctx, index, &driver.Document{
|
||||
DocID: id,
|
||||
Content: content,
|
||||
Metadata: metadata,
|
||||
})
|
||||
}
|
||||
|
||||
// Map convert the assistant to a map
|
||||
func (ast *Assistant) Map() map[string]interface{} {
|
||||
|
||||
if ast == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"assistant_id": ast.ID,
|
||||
"type": ast.Type,
|
||||
"name": ast.Name,
|
||||
"readonly": ast.Readonly,
|
||||
"avatar": ast.Avatar,
|
||||
"connector": ast.Connector,
|
||||
"path": ast.Path,
|
||||
"built_in": ast.BuiltIn,
|
||||
"sort": ast.Sort,
|
||||
"description": ast.Description,
|
||||
"options": ast.Options,
|
||||
"prompts": ast.Prompts,
|
||||
"tags": ast.Tags,
|
||||
"mentionable": ast.Mentionable,
|
||||
"automated": ast.Automated,
|
||||
"created_at": timeToMySQLFormat(ast.CreatedAt),
|
||||
"updated_at": timeToMySQLFormat(ast.UpdatedAt),
|
||||
}
|
||||
}
|
||||
|
||||
// Validate validates the assistant configuration
|
||||
func (ast *Assistant) Validate() error {
|
||||
if ast.ID == "" {
|
||||
return fmt.Errorf("assistant_id is required")
|
||||
}
|
||||
if ast.Name == "" {
|
||||
return fmt.Errorf("name is required")
|
||||
}
|
||||
if ast.Connector == "" {
|
||||
return fmt.Errorf("connector is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Clone creates a deep copy of the assistant
|
||||
func (ast *Assistant) Clone() *Assistant {
|
||||
if ast == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
clone := &Assistant{
|
||||
ID: ast.ID,
|
||||
Type: ast.Type,
|
||||
Name: ast.Name,
|
||||
Avatar: ast.Avatar,
|
||||
Connector: ast.Connector,
|
||||
Path: ast.Path,
|
||||
BuiltIn: ast.BuiltIn,
|
||||
Sort: ast.Sort,
|
||||
Description: ast.Description,
|
||||
Readonly: ast.Readonly,
|
||||
Mentionable: ast.Mentionable,
|
||||
Automated: ast.Automated,
|
||||
Script: ast.Script,
|
||||
API: ast.API,
|
||||
}
|
||||
|
||||
// Deep copy tags
|
||||
if ast.Tags != nil {
|
||||
clone.Tags = make([]string, len(ast.Tags))
|
||||
copy(clone.Tags, ast.Tags)
|
||||
}
|
||||
|
||||
// Deep copy options
|
||||
if ast.Options != nil {
|
||||
clone.Options = make(map[string]interface{})
|
||||
for k, v := range ast.Options {
|
||||
clone.Options[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
// Deep copy prompts
|
||||
if ast.Prompts != nil {
|
||||
clone.Prompts = make([]Prompt, len(ast.Prompts))
|
||||
copy(clone.Prompts, ast.Prompts)
|
||||
}
|
||||
|
||||
// Deep copy flows
|
||||
if ast.Flows != nil {
|
||||
clone.Flows = make([]map[string]interface{}, len(ast.Flows))
|
||||
for i, flow := range ast.Flows {
|
||||
cloneFlow := make(map[string]interface{})
|
||||
for k, v := range flow {
|
||||
cloneFlow[k] = v
|
||||
}
|
||||
clone.Flows[i] = cloneFlow
|
||||
}
|
||||
}
|
||||
|
||||
return clone
|
||||
}
|
||||
|
||||
// Update updates the assistant properties
|
||||
func (ast *Assistant) Update(data map[string]interface{}) error {
|
||||
if ast == nil {
|
||||
return fmt.Errorf("assistant is nil")
|
||||
}
|
||||
|
||||
if v, ok := data["name"].(string); ok {
|
||||
ast.Name = v
|
||||
}
|
||||
if v, ok := data["avatar"].(string); ok {
|
||||
ast.Avatar = v
|
||||
}
|
||||
if v, ok := data["description"].(string); ok {
|
||||
ast.Description = v
|
||||
}
|
||||
if v, ok := data["connector"].(string); ok {
|
||||
ast.Connector = v
|
||||
}
|
||||
if v, ok := data["type"].(string); ok {
|
||||
ast.Type = v
|
||||
}
|
||||
if v, ok := data["sort"].(int); ok {
|
||||
ast.Sort = v
|
||||
}
|
||||
if v, ok := data["mentionable"].(bool); ok {
|
||||
ast.Mentionable = v
|
||||
}
|
||||
if v, ok := data["automated"].(bool); ok {
|
||||
ast.Automated = v
|
||||
}
|
||||
if v, ok := data["tags"].([]string); ok {
|
||||
ast.Tags = v
|
||||
}
|
||||
if v, ok := data["options"].(map[string]interface{}); ok {
|
||||
ast.Options = v
|
||||
}
|
||||
|
||||
return ast.Validate()
|
||||
}
|
||||
|
|
@ -67,16 +67,15 @@ func LoadBuiltIn() error {
|
|||
assistant.Tags = []string{"Built-in"}
|
||||
}
|
||||
|
||||
// Save the assistant
|
||||
err = assistant.Save()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
sort++
|
||||
loaded.Put(assistant)
|
||||
|
||||
// Save the assistant
|
||||
if storage != nil {
|
||||
_, err := storage.SaveAssistant(assistant.Map())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
|
|
@ -91,11 +90,12 @@ func SetStorage(s store.Store) {
|
|||
// e: the RAG engine
|
||||
// u: the RAG file uploader
|
||||
// v: the RAG vectorizer
|
||||
func SetRAG(e driver.Engine, u driver.FileUpload, v driver.Vectorizer) {
|
||||
func SetRAG(e driver.Engine, u driver.FileUpload, v driver.Vectorizer, setting RAGSetting) {
|
||||
rag = &RAG{
|
||||
Engine: e,
|
||||
Uploader: u,
|
||||
Vectorizer: v,
|
||||
Setting: setting,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -177,24 +177,30 @@ func LoadPath(path string) (*Assistant, error) {
|
|||
data["assistant_id"] = id
|
||||
data["type"] = "assistant"
|
||||
data["path"] = path
|
||||
|
||||
updatedAt := int64(0)
|
||||
|
||||
// prompts
|
||||
promptsfile := filepath.Join(path, "prompts.yml")
|
||||
if has, _ := app.Exists(promptsfile); has {
|
||||
prompts, err := loadPrompts(promptsfile, path)
|
||||
prompts, ts, err := loadPrompts(promptsfile, path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
data["prompts"] = prompts
|
||||
data["updated_at"] = ts
|
||||
updatedAt = ts
|
||||
}
|
||||
|
||||
// load script
|
||||
scriptfile := filepath.Join(path, "src", "index.ts")
|
||||
if has, _ := app.Exists(scriptfile); has {
|
||||
script, err := loadScript(scriptfile, path)
|
||||
script, ts, err := loadScript(scriptfile, path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
data["script"] = script
|
||||
data["updated_at"] = max(updatedAt, ts)
|
||||
}
|
||||
|
||||
// load functions
|
||||
|
|
@ -307,19 +313,42 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
|
|||
}
|
||||
}
|
||||
|
||||
// created_at
|
||||
if v, has := data["created_at"]; has {
|
||||
ts, err := getTimestamp(v)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
assistant.CreatedAt = ts
|
||||
}
|
||||
|
||||
// updated_at
|
||||
if v, has := data["updated_at"]; has {
|
||||
ts, err := getTimestamp(v)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
assistant.UpdatedAt = ts
|
||||
}
|
||||
|
||||
return assistant, nil
|
||||
}
|
||||
|
||||
func loadPrompts(file string, root string) (string, error) {
|
||||
func loadPrompts(file string, root string) (string, int64, error) {
|
||||
|
||||
app, err := fs.Get("app")
|
||||
if err != nil {
|
||||
return "", err
|
||||
return "", 0, err
|
||||
}
|
||||
|
||||
ts, err := app.ModTime(file)
|
||||
if err != nil {
|
||||
return "", 0, err
|
||||
}
|
||||
|
||||
prompts, err := app.ReadFile(file)
|
||||
if err != nil {
|
||||
return "", err
|
||||
return "", 0, err
|
||||
}
|
||||
|
||||
re := regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`)
|
||||
|
|
@ -339,11 +368,27 @@ func loadPrompts(file string, root string) (string, error) {
|
|||
return []byte(formattedContent)
|
||||
})
|
||||
|
||||
return string(prompts), nil
|
||||
return string(prompts), ts.UnixNano(), nil
|
||||
}
|
||||
|
||||
func loadScript(file string, root string) (*v8.Script, error) {
|
||||
return v8.Load(file, share.ID(root, file))
|
||||
func loadScript(file string, root string) (*v8.Script, int64, error) {
|
||||
|
||||
app, err := fs.Get("app")
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
ts, err := app.ModTime(file)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
script, err := v8.Load(file, share.ID(root, file))
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
|
||||
return script, ts.UnixNano(), nil
|
||||
}
|
||||
|
||||
func loadScriptSource(source string, file string) (*v8.Script, error) {
|
||||
|
|
@ -353,151 +398,3 @@ func loadScriptSource(source string, file string) (*v8.Script, error) {
|
|||
}
|
||||
return script, nil
|
||||
}
|
||||
|
||||
// Save save the assistant
|
||||
func (ast *Assistant) Save() error {
|
||||
if storage == nil {
|
||||
return fmt.Errorf("storage is not set")
|
||||
}
|
||||
|
||||
_, err := storage.SaveAssistant(ast.Map())
|
||||
return err
|
||||
}
|
||||
|
||||
// Map convert the assistant to a map
|
||||
func (ast *Assistant) Map() map[string]interface{} {
|
||||
|
||||
if ast == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"assistant_id": ast.ID,
|
||||
"type": ast.Type,
|
||||
"name": ast.Name,
|
||||
"readonly": ast.Readonly,
|
||||
"avatar": ast.Avatar,
|
||||
"connector": ast.Connector,
|
||||
"path": ast.Path,
|
||||
"built_in": ast.BuiltIn,
|
||||
"sort": ast.Sort,
|
||||
"description": ast.Description,
|
||||
"options": ast.Options,
|
||||
"prompts": ast.Prompts,
|
||||
"tags": ast.Tags,
|
||||
"mentionable": ast.Mentionable,
|
||||
"automated": ast.Automated,
|
||||
}
|
||||
}
|
||||
|
||||
// Validate validates the assistant configuration
|
||||
func (ast *Assistant) Validate() error {
|
||||
if ast.ID == "" {
|
||||
return fmt.Errorf("assistant_id is required")
|
||||
}
|
||||
if ast.Name == "" {
|
||||
return fmt.Errorf("name is required")
|
||||
}
|
||||
if ast.Connector == "" {
|
||||
return fmt.Errorf("connector is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Clone creates a deep copy of the assistant
|
||||
func (ast *Assistant) Clone() *Assistant {
|
||||
if ast == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
clone := &Assistant{
|
||||
ID: ast.ID,
|
||||
Type: ast.Type,
|
||||
Name: ast.Name,
|
||||
Avatar: ast.Avatar,
|
||||
Connector: ast.Connector,
|
||||
Path: ast.Path,
|
||||
BuiltIn: ast.BuiltIn,
|
||||
Sort: ast.Sort,
|
||||
Description: ast.Description,
|
||||
Readonly: ast.Readonly,
|
||||
Mentionable: ast.Mentionable,
|
||||
Automated: ast.Automated,
|
||||
Script: ast.Script,
|
||||
API: ast.API,
|
||||
}
|
||||
|
||||
// Deep copy tags
|
||||
if ast.Tags != nil {
|
||||
clone.Tags = make([]string, len(ast.Tags))
|
||||
copy(clone.Tags, ast.Tags)
|
||||
}
|
||||
|
||||
// Deep copy options
|
||||
if ast.Options != nil {
|
||||
clone.Options = make(map[string]interface{})
|
||||
for k, v := range ast.Options {
|
||||
clone.Options[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
// Deep copy prompts
|
||||
if ast.Prompts != nil {
|
||||
clone.Prompts = make([]Prompt, len(ast.Prompts))
|
||||
copy(clone.Prompts, ast.Prompts)
|
||||
}
|
||||
|
||||
// Deep copy flows
|
||||
if ast.Flows != nil {
|
||||
clone.Flows = make([]map[string]interface{}, len(ast.Flows))
|
||||
for i, flow := range ast.Flows {
|
||||
cloneFlow := make(map[string]interface{})
|
||||
for k, v := range flow {
|
||||
cloneFlow[k] = v
|
||||
}
|
||||
clone.Flows[i] = cloneFlow
|
||||
}
|
||||
}
|
||||
|
||||
return clone
|
||||
}
|
||||
|
||||
// Update updates the assistant properties
|
||||
func (ast *Assistant) Update(data map[string]interface{}) error {
|
||||
if ast == nil {
|
||||
return fmt.Errorf("assistant is nil")
|
||||
}
|
||||
|
||||
if v, ok := data["name"].(string); ok {
|
||||
ast.Name = v
|
||||
}
|
||||
if v, ok := data["avatar"].(string); ok {
|
||||
ast.Avatar = v
|
||||
}
|
||||
if v, ok := data["description"].(string); ok {
|
||||
ast.Description = v
|
||||
}
|
||||
if v, ok := data["connector"].(string); ok {
|
||||
ast.Connector = v
|
||||
}
|
||||
if v, ok := data["type"].(string); ok {
|
||||
ast.Type = v
|
||||
}
|
||||
if v, ok := data["sort"].(int); ok {
|
||||
ast.Sort = v
|
||||
}
|
||||
if v, ok := data["mentionable"].(bool); ok {
|
||||
ast.Mentionable = v
|
||||
}
|
||||
if v, ok := data["automated"].(bool); ok {
|
||||
ast.Automated = v
|
||||
}
|
||||
if v, ok := data["tags"].([]string); ok {
|
||||
ast.Tags = v
|
||||
}
|
||||
if v, ok := data["options"].(map[string]interface{}); ok {
|
||||
ast.Options = v
|
||||
}
|
||||
|
||||
return ast.Validate()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -22,6 +22,12 @@ type RAG struct {
|
|||
Engine driver.Engine
|
||||
Uploader driver.FileUpload
|
||||
Vectorizer driver.Vectorizer
|
||||
Setting RAGSetting
|
||||
}
|
||||
|
||||
// RAGSetting the RAG setting
|
||||
type RAGSetting struct {
|
||||
IndexPrefix string `json:"index_prefix" yaml:"index_prefix"`
|
||||
}
|
||||
|
||||
// Prompt a prompt
|
||||
|
|
@ -59,6 +65,8 @@ type Assistant struct {
|
|||
Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows
|
||||
Script *v8.Script `json:"-" yaml:"-"` // Assistant Script
|
||||
API API `json:"-" yaml:"-"` // Assistant API
|
||||
CreatedAt int64 `json:"created_at"` // Creation timestamp
|
||||
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
|
||||
}
|
||||
|
||||
// File the file
|
||||
|
|
|
|||
44
neo/assistant/utils.go
Normal file
44
neo/assistant/utils.go
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
package assistant
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
func getTimestamp(v interface{}) (int64, error) {
|
||||
switch v := v.(type) {
|
||||
case int64:
|
||||
return v, nil
|
||||
case int:
|
||||
return int64(v), nil
|
||||
|
||||
case string:
|
||||
if ts, err := time.Parse(time.RFC3339, v); err == nil {
|
||||
return ts.UnixNano(), nil
|
||||
}
|
||||
|
||||
// MySQL format
|
||||
if ts, err := time.Parse("2006-01-02 15:04:05", v); err == nil {
|
||||
return ts.UnixNano(), nil
|
||||
}
|
||||
|
||||
// UnixNano format
|
||||
if ts, err := strconv.ParseInt(v, 10, 64); err == nil {
|
||||
return ts, nil
|
||||
}
|
||||
|
||||
}
|
||||
return 0, fmt.Errorf("invalid timestamp type")
|
||||
}
|
||||
|
||||
func stringToTimestamp(v string) (int64, error) {
|
||||
return strconv.ParseInt(v, 10, 64)
|
||||
}
|
||||
|
||||
func timeToMySQLFormat(ts int64) string {
|
||||
if ts == 0 {
|
||||
return "0000-00-00 00:00:00"
|
||||
}
|
||||
return time.Unix(ts/1e9, ts%1e9).Format("2006-01-02 15:04:05")
|
||||
}
|
||||
|
|
@ -119,7 +119,14 @@ func (neo *DSL) initAssistant() error {
|
|||
|
||||
// Assistant RAG
|
||||
if Neo.RAG != nil {
|
||||
assistant.SetRAG(Neo.RAG.Engine(), Neo.RAG.FileUpload(), Neo.RAG.Vectorizer())
|
||||
assistant.SetRAG(
|
||||
Neo.RAG.Engine(),
|
||||
Neo.RAG.FileUpload(),
|
||||
Neo.RAG.Vectorizer(),
|
||||
assistant.RAGSetting{
|
||||
IndexPrefix: Neo.RAGSetting.IndexPrefix,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// Load Built-in Assistants
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue