diff --git a/example/webui/agentpool.go b/example/webui/agentpool.go new file mode 100644 index 0000000..b0ecb73 --- /dev/null +++ b/example/webui/agentpool.go @@ -0,0 +1,148 @@ +package main + +import ( + "encoding/json" + "fmt" + "os" + + . "github.com/mudler/local-agent-framework/agent" +) + +type AgentConfig struct { +} + +type AgentPool struct { + file string + pool AgentPoolData + agents map[string]*Agent +} + +type AgentPoolData map[string]AgentConfig + +func NewAgentPool(file string) (*AgentPool, error) { + // if file exists, try to load an existing pool. + // if file does not exist, create a new pool. + + if _, err := os.Stat(file); err != nil { + // file does not exist, create a new pool + return &AgentPool{ + file: file, + agents: make(map[string]*Agent), + pool: make(map[string]AgentConfig), + }, nil + } + + poolData, err := loadPoolFromFile(file) + if err != nil { + return nil, err + } + return &AgentPool{ + file: file, + agents: make(map[string]*Agent), + pool: *poolData, + }, nil +} + +// CreateAgent adds a new agent to the pool +// and starts it. +// It also saves the state to the file. +func (a *AgentPool) CreateAgent(name string, agentConfig *AgentConfig) error { + if _, ok := a.pool[name]; ok { + return fmt.Errorf("agent %s already exists", name) + } + a.pool[name] = *agentConfig + if err := a.Save(); err != nil { + return err + } + + return a.startAgentWithConfig(name, agentConfig) +} + +func (a *AgentPool) startAgentWithConfig(name string, config *AgentConfig) error { + + agent, err := New( + WithModel("hermes-2-pro-mistral"), + ) + if err != nil { + return err + } + + a.agents[name] = agent + + go func() { + if err := agent.Run(); err != nil { + panic(err) + } + }() + + return nil +} + +// Starts all the agents in the pool +func (a *AgentPool) StartAll() error { + for name, config := range a.pool { + if a.agents[name] != nil { // Agent already started + continue + } + if err := a.startAgentWithConfig(name, &config); err != nil { + return err + } + } + return nil +} + +func (a *AgentPool) StopAll() { + for _, agent := range a.agents { + agent.Stop() + } +} + +func (a *AgentPool) Stop(name string) { + if agent, ok := a.agents[name]; ok { + agent.Stop() + } +} + +func (a *AgentPool) Start(name string) error { + if agent, ok := a.agents[name]; ok { + return agent.Run() + } + if config, ok := a.pool[name]; ok { + return a.startAgentWithConfig(name, &config) + } + + return fmt.Errorf("agent %s not found", name) +} + +func (a *AgentPool) Remove(name string) error { + a.Stop(name) + delete(a.agents, name) + delete(a.pool, name) + if err := a.Save(); err != nil { + return err + } + return nil +} + +func (a *AgentPool) Save() error { + data, err := json.Marshal(a.pool) + if err != nil { + return err + } + return os.WriteFile(a.file, data, 0644) +} + +func (a *AgentPool) GetAgent(name string) *Agent { + return a.agents[name] +} + +func loadPoolFromFile(path string) (*AgentPoolData, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + poolData := &AgentPoolData{} + err = json.Unmarshal(data, poolData) + return poolData, err +}