fdf71f4cfd
- Bump Go version from 1.22.2 to 1.23 for improved performance. - Update dependencies in go.mod and go.sum, including: - github.com/PuerkitoBio/goquery to v1.10.1 - github.com/dchest/captcha to v1.1.0 - github.com/evanw/esbuild to v0.24.2 - github.com/fatih/color to v1.18.0 - github.com/fsnotify/fsnotify to v1.8.0 - github.com/spf13/cobra to v1.8.1 - github.com/stretchr/testify to v1.10.0 - github.com/xuri/excelize/v2 to v2.9.0 - Introduce RAG (Retrieval-Augmented Generation) initialization in the Neo API, enhancing assistant capabilities. - Refactor workflows to use Go 1.23 for testing and CI processes, ensuring compatibility with the latest features.
120 lines
2.8 KiB
Go
120 lines
2.8 KiB
Go
package rag
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/yaoapp/gou/rag"
|
|
"github.com/yaoapp/gou/rag/driver"
|
|
)
|
|
|
|
// RAG the RAG instance
|
|
type RAG struct {
|
|
setting Setting
|
|
engine driver.Engine
|
|
vectorizer driver.Vectorizer
|
|
fileUpload driver.FileUpload
|
|
}
|
|
|
|
// 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]string {
|
|
converted := make(map[string]string)
|
|
for k, v := range options {
|
|
if str, ok := v.(string); ok {
|
|
converted[k] = parseEnvValue(str)
|
|
}
|
|
}
|
|
return converted
|
|
}
|
|
|
|
// New create a new RAG instance
|
|
func New(setting Setting) (*RAG, error) {
|
|
if setting.Engine.Driver == "" {
|
|
return nil, fmt.Errorf("engine driver is required")
|
|
}
|
|
|
|
if setting.Vectorizer.Driver == "" {
|
|
return nil, fmt.Errorf("vectorizer driver is required")
|
|
}
|
|
|
|
// Set default values
|
|
if setting.Upload.ChunkSize == 0 {
|
|
setting.Upload.ChunkSize = 1024
|
|
}
|
|
|
|
if setting.Upload.ChunkOverlap == 0 {
|
|
setting.Upload.ChunkOverlap = 256
|
|
}
|
|
|
|
if setting.IndexPrefix == "" {
|
|
setting.IndexPrefix = "yao_neo_"
|
|
}
|
|
|
|
// Convert options map for vectorizer and handle environment variables
|
|
vectorizerOpts := convertOptions(setting.Vectorizer.Options)
|
|
|
|
// Create vectorizer
|
|
vectorizer, err := rag.NewVectorizer(setting.Vectorizer.Driver, driver.VectorizeConfig{
|
|
Model: vectorizerOpts["model"],
|
|
Options: vectorizerOpts,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create vectorizer: %v", err)
|
|
}
|
|
|
|
// Convert options map for engine and handle environment variables
|
|
engineOpts := convertOptions(setting.Engine.Options)
|
|
|
|
// Create engine
|
|
engine, err := rag.NewEngine(setting.Engine.Driver, driver.IndexConfig{
|
|
Options: engineOpts,
|
|
}, vectorizer)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create engine: %v", err)
|
|
}
|
|
|
|
// Create file upload
|
|
fileUpload, err := rag.NewFileUpload(setting.Engine.Driver, engine, vectorizer)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("create file upload: %v", err)
|
|
}
|
|
|
|
return &RAG{
|
|
setting: setting,
|
|
engine: engine,
|
|
vectorizer: vectorizer,
|
|
fileUpload: fileUpload,
|
|
}, nil
|
|
}
|
|
|
|
// Setting get the RAG settings
|
|
func (rag *RAG) Setting() Setting {
|
|
return rag.setting
|
|
}
|
|
|
|
// Engine get the vector database engine
|
|
func (rag *RAG) Engine() driver.Engine {
|
|
return rag.engine
|
|
}
|
|
|
|
// Vectorizer get the text vectorizer
|
|
func (rag *RAG) Vectorizer() driver.Vectorizer {
|
|
return rag.vectorizer
|
|
}
|
|
|
|
// FileUpload get the file upload handler
|
|
func (rag *RAG) FileUpload() driver.FileUpload {
|
|
return rag.fileUpload
|
|
}
|