Enhance assistant loading and ID generation

- Update loadMap function to handle prompts as both string and array types, improving flexibility in data loading.
- Refactor assistant ID generation in SaveAssistant method to use a dedicated GenerateAssistantID function, ensuring unique IDs are created with error handling.
- Add GenerateAssistantID method to create a random 6-digit ID while checking for uniqueness in the database.
This commit is contained in:
Max 2025-03-17 18:28:02 +08:00
parent 38b910108b
commit 586ad9866c
3 changed files with 65 additions and 8 deletions

View file

@ -424,13 +424,33 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
} }
// prompts // prompts
if v, ok := data["prompts"].(string); ok { if prompts, has := data["prompts"]; has {
var prompts []Prompt
err := yaml.Unmarshal([]byte(v), &prompts) switch v := prompts.(type) {
if err != nil { case []Prompt:
return nil, err assistant.Prompts = v
case string:
var prompts []Prompt
err := yaml.Unmarshal([]byte(v), &prompts)
if err != nil {
return nil, err
}
assistant.Prompts = prompts
default:
raw, err := jsoniter.Marshal(v)
if err != nil {
return nil, err
}
var prompts []Prompt
err = jsoniter.Unmarshal(raw, &prompts)
if err != nil {
return nil, err
}
assistant.Prompts = prompts
} }
assistant.Prompts = prompts
} }
// tools // tools

View file

@ -33,8 +33,11 @@ func getTimestamp(v interface{}) (int64, error) {
return ts, nil return ts, nil
} }
case nil:
return 0, nil
} }
return 0, fmt.Errorf("invalid timestamp type")
return 0, fmt.Errorf("invalid timestamp type %T", v)
} }
func stringToTimestamp(v string) (int64, error) { func stringToTimestamp(v string) (int64, error) {

View file

@ -963,7 +963,11 @@ func (conv *Xun) SaveAssistant(assistant map[string]interface{}) (interface{}, e
// Generate assistant_id if not provided // Generate assistant_id if not provided
if _, ok := assistantCopy["assistant_id"]; !ok { if _, ok := assistantCopy["assistant_id"]; !ok {
assistantCopy["assistant_id"] = uuid.New().String() var err error
assistantCopy["assistant_id"], err = conv.GenerateAssistantID()
if err != nil {
return nil, err
}
} }
// Check if assistant exists // Check if assistant exists
@ -1359,3 +1363,33 @@ func (conv *Xun) GetHistoryWithFilter(sid string, cid string, filter ChatFilter)
return res, nil return res, nil
} }
// GenerateAssistantID generates a random-looking 6-digit ID
func (conv *Xun) GenerateAssistantID() (string, error) {
maxAttempts := 10 // Maximum number of attempts to generate a unique ID
for i := 0; i < maxAttempts; i++ {
// Generate a random number using timestamp and some bit operations
timestamp := time.Now().UnixNano()
random := (timestamp ^ (timestamp >> 12)) % 1000000
hash := fmt.Sprintf("%06d", random)
// Check if this ID already exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", hash).
Exists()
if err != nil {
return "", err
}
if !exists {
return hash, nil
}
// If ID exists, wait a bit and try again
time.Sleep(time.Millisecond)
}
return "", fmt.Errorf("failed to generate unique ID after %d attempts", maxAttempts)
}