8fdd1a3a6c
- Introduced a new `workspace` field across various robot-related structures, including `CreateRobotRequest`, `UpdateRobotRequest`, and `RobotResponse`, allowing for better organization and management of robots within specific workspaces. - Updated database queries and response mappings to accommodate the new workspace field, ensuring seamless integration with existing functionalities. - Enhanced agent execution context to include workspace information, improving the contextual awareness of agents during operations. - Added tests to validate the creation and updating of robots with workspace data, ensuring robust functionality and backward compatibility.
121 lines
2.7 KiB
Go
121 lines
2.7 KiB
Go
package cache
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"github.com/yaoapp/gou/model"
|
|
"github.com/yaoapp/kun/maps"
|
|
"github.com/yaoapp/yao/agent/robot/types"
|
|
)
|
|
|
|
// memberModel is the model name for member table
|
|
// Can be changed via SetMemberModel() during system initialization
|
|
var memberModel = "__yao.member"
|
|
|
|
// memberFields are the fields to select when loading robots
|
|
var memberFields = []interface{}{
|
|
"id",
|
|
"member_id",
|
|
"team_id",
|
|
"display_name",
|
|
"bio",
|
|
"system_prompt",
|
|
"robot_status",
|
|
"autonomous_mode",
|
|
"robot_config",
|
|
"robot_email",
|
|
"agents",
|
|
"mcp_servers",
|
|
"manager_id",
|
|
"language_model",
|
|
"workspace",
|
|
}
|
|
|
|
// SetMemberModel sets the member model name
|
|
// Call this during system initialization to override the default
|
|
func SetMemberModel(model string) {
|
|
if model != "" {
|
|
memberModel = model
|
|
}
|
|
}
|
|
|
|
// Load loads all active robots from database with pagination
|
|
// Query: member_type='robot' AND status='active'
|
|
func (c *Cache) Load(ctx *types.Context) error {
|
|
m := model.Select(memberModel)
|
|
|
|
// Clear existing cache first
|
|
c.mu.Lock()
|
|
c.robots = make(map[string]*types.Robot)
|
|
c.byTeam = make(map[string][]string)
|
|
c.mu.Unlock()
|
|
|
|
// Paginate to handle large number of robots
|
|
page := 1
|
|
pageSize := 100 // load 100 robots per page
|
|
totalLoaded := 0
|
|
|
|
for {
|
|
// Query with pagination
|
|
result, err := m.Paginate(model.QueryParam{
|
|
Select: memberFields,
|
|
Wheres: []model.QueryWhere{
|
|
{Column: "member_type", Value: "robot"},
|
|
{Column: "status", Value: "active"},
|
|
},
|
|
}, page, pageSize)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to load robots (page %d): %w", page, err)
|
|
}
|
|
|
|
// Extract records from pagination result
|
|
data, ok := result.Get("data").([]maps.MapStr)
|
|
if !ok || len(data) == 0 {
|
|
break
|
|
}
|
|
|
|
// Parse and add each robot
|
|
for _, record := range data {
|
|
robot, err := types.NewRobotFromMap(map[string]interface{}(record))
|
|
if err != nil {
|
|
// Log error but continue loading other robots
|
|
continue
|
|
}
|
|
c.Add(robot)
|
|
totalLoaded++
|
|
}
|
|
|
|
// Check if there are more pages
|
|
total, _ := result.Get("total").(int)
|
|
if totalLoaded >= total {
|
|
break
|
|
}
|
|
|
|
page++
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// LoadByID loads a single robot from database by member ID
|
|
func (c *Cache) LoadByID(ctx *types.Context, memberID string) (*types.Robot, error) {
|
|
m := model.Select(memberModel)
|
|
|
|
records, err := m.Get(model.QueryParam{
|
|
Select: memberFields,
|
|
Wheres: []model.QueryWhere{
|
|
{Column: "member_id", Value: memberID},
|
|
{Column: "member_type", Value: "robot"},
|
|
},
|
|
Limit: 1,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to load robot %s: %w", memberID, err)
|
|
}
|
|
|
|
if len(records) == 0 {
|
|
return nil, types.ErrRobotNotFound
|
|
}
|
|
|
|
return types.NewRobotFromMap(map[string]interface{}(records[0]))
|
|
}
|