diff --git a/pkg/agents/console.go b/pkg/agents/console.go index 87dd580..a519b7c 100644 --- a/pkg/agents/console.go +++ b/pkg/agents/console.go @@ -2,24 +2,284 @@ package agents import ( "context" + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" + + gollm "gitea.stevedudenhoeffer.com/steve/go-llm" + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/mount" + "github.com/docker/docker/client" "gitea.stevedudenhoeffer.com/steve/answer/pkg/agents/console" + "gitea.stevedudenhoeffer.com/steve/answer/pkg/agents/shared" ) -// Console will construct a console.Agent, execute the agent, and return the knowledge gained, the output directory, -// and any possible error. -func (a Agent) Console(ctx context.Context, questions []string) (Knowledge, string, error) { - con := console.Agent{ - Agent: a, - Model: a.model, - ContextualInformation: a.contextualInformation, - MaxCommands: 50, - } - - res, err := con.Answer(ctx, questions) - if err != nil { - return Knowledge{}, "", err - } - - return res.Knowledge, res.Directory, nil +type ConsoleConfig struct { + // MaxCommands is how many total commands this console can run. + MaxCommands int `json:"max_commands"` + OnCommandStart func(ctx context.Context, command string) error + OnCommandDone func(ctx context.Context, command string, output string, err error) error +} + +func (a Agent) Console(ctx context.Context, questions []string, cfg ConsoleConfig) (Knowledge, string, error) { + + if cfg.MaxCommands <= 0 { + cfg.MaxCommands = 10000 + } + + var resKnowledge = Knowledge{ + OriginalQuestions: questions, + RemainingQuestions: questions, + } + + // create a temporary scratch directory + dir, err := os.MkdirTemp("", "console-") + if err != nil { + return resKnowledge, "", err + } + + var resDirectory = dir + + cl, err := client.NewClientWithOpts(client.FromEnv) + if err != nil { + return resKnowledge, resDirectory, err + } + defer cl.Close() + + mounts := []mount.Mount{ + { + Type: mount.TypeBind, + Source: dir, + Target: "/home/user", + }, + } + + c, err := console.CreateContainer(ctx, cl, console.ContainerConfig{ + Config: &container.Config{ + Image: "ubuntu:latest", + Cmd: []string{"tail", "-f", "/dev/null"}, + Tty: true, + WorkingDir: "/home/user", + }, + HostConfig: &container.HostConfig{ + AutoRemove: true, + Mounts: mounts, + }, + Name: filepath.Base(dir), + }) + + if err != nil { + return resKnowledge, resDirectory, fmt.Errorf("failed to create container: %w", err) + } + defer func() { + _ = c.Close(ctx) + }() + + slog.Info("starting container", "dir", dir, "container", fmt.Sprintf("%+v", c)) + err = c.Start(ctx) + if err != nil { + return resKnowledge, resDirectory, fmt.Errorf("failed to start container: %w", err) + } + + // Run the model + + var history console.Executions + var keepGoing = true + + opwd, epwd := c.Execute(ctx, "ls -al /home") + + fmt.Println(opwd) + slog.Info("pwd", "pwd", opwd, "epwd", epwd) + + tools := map[string]gollm.Function{ + "exit": gollm.NewFunction( + "exit", + "exit the container", + func(ctx *gollm.Context, args struct { + RemainingQuestions []string `description:"any remaining questions that remain unanswered"` + }) (any, error) { + keepGoing = false + return "exiting", nil + }), + + "write": gollm.NewFunction( + "write", + "write a file in the /root directory", + func(ctx *gollm.Context, args struct { + Filename string `description:"The name of the file to write"` + Content string `description:"The content of the file to write"` + }) (any, error) { + target, err := console.SafeJoinPath(dir, args.Filename) + if err != nil { + return "", err + } + + f, err := os.Create(target) + if err != nil { + return "", err + } + defer f.Close() + + _, err = f.WriteString(args.Content) + if err != nil { + return "", err + } + + return "wrote file", nil + }), + + "read": gollm.NewFunction( + "read", + "read a file in the /root directory", + func(ctx *gollm.Context, args struct { + Filename string `description:"The name of the file to read"` + }) (any, error) { + target, err := console.SafeJoinPath(dir, args.Filename) + if err != nil { + return "", err + } + + b, err := os.ReadFile(target) + if err != nil { + return "", err + } + + return string(b), nil + }), + + "execute": gollm.NewFunction( + "execute", + "execute a command in the container", + func(ctx *gollm.Context, args struct { + Command string `description:"The command to execute"` + }) (any, error) { + if len(history) >= cfg.MaxCommands { + return "too many commands", nil + } + + if cfg.OnCommandStart != nil { + err := cfg.OnCommandStart(ctx, args.Command) + if err != nil { + return "", err + } + } + + var res string + // if the command starts with sudo then we need to use the sudo function + if strings.HasPrefix(args.Command, "sudo ") { + res, err = c.Sudo(ctx, args.Command[5:]) + } else { + res, err = c.Execute(ctx, args.Command) + } + + if cfg.OnCommandDone != nil { + err = cfg.OnCommandDone(ctx, args.Command, res, nil) + if err != nil { + return "", err + } + } + + history = append(history, console.Execution{ + Command: args.Command, + Output: res, + }) + + return res, nil + }), + + "sudo": gollm.NewFunction( + "sudo", + "execute a command in the container", + func(ctx *gollm.Context, args struct { + Command string `description:"The command to execute"` + }) (any, error) { + if len(history) >= cfg.MaxCommands { + return "too many commands", nil + } + + if cfg.OnCommandStart != nil { + err := cfg.OnCommandStart(ctx, args.Command) + if err != nil { + return "", err + } + } + + res, err := c.Sudo(ctx, args.Command) + + if cfg.OnCommandDone != nil { + err = cfg.OnCommandDone(ctx, args.Command, res, nil) + if err != nil { + return "", err + } + } + + if err != nil { + res = "error executing: " + err.Error() + } + + history = append(history, console.Execution{ + Command: "sudo " + args.Command, + Output: res, + }) + + return res, nil + }), + } + + for i := 0; i < cfg.MaxCommands && len(history) < cfg.MaxCommands && keepGoing; i++ { + + systemPrompt := `You are now in a shell in a container of the ubuntu:latest image to answer a question asked by the user, it is very basic install of ubuntu, simple things (like python) are not preinstalled but can be installed via apt. You will be run multiple times and gain knowledge throughout the process.` + + if len(history) < cfg.MaxCommands { + systemPrompt += `You can run any command you like to get to the needed results.` + } + + systemPrompt += `Alternatively, you can use the tool "write" to write a file in the home directory, and also the tool "read" to read a file in the home directory. +When you are done, please use "exit" to exit the container. +Respond with any number of commands to answer the question, they will be executed in order.` + + var toolbox []gollm.Function + + // add unrestricted tools + toolbox = append(toolbox, tools["exit"], tools["write"], tools["read"]) + + if len(history) < cfg.MaxCommands { + toolbox = append(toolbox, tools["execute"], tools["sudo"]) + } + + kw := shared.KnowledgeWorker{ + Model: a.model, + ToolBox: gollm.NewToolBox(toolbox...), + ContextualInformation: a.contextualInformation, + OnNewFunction: func(ctx context.Context, funcName string, args string) (any, error) { + slog.Info("new function called", "function name", funcName, "args", args) + return nil, nil + }, + } + + slog.Info("answering question", "question", questions[0]) + r, err := kw.Answer(ctx, &resKnowledge, systemPrompt, "", "", history.ToGeneralButLastMessageHistory(), func(res gollm.ToolCallResponse) { + }) + + if err != nil { + return resKnowledge, resDirectory, fmt.Errorf("error answering question: %w", err) + } + + if len(r.Knowledge) > 0 { + slog.Info("answered question and learned", "knowledge", r.Knowledge) + } else { + slog.Info("answered question and learned nothing") + } + + resKnowledge, err = a.KnowledgeIntegrate(ctx, resKnowledge, r) + if err != nil { + return resKnowledge, resDirectory, fmt.Errorf("error integrating knowledge: %w", err) + } + + slog.Info("knowledge integrated", "question", questions[0], "knowledge", resKnowledge) + } + return resKnowledge, resDirectory, nil } diff --git a/pkg/agents/console/agent.go b/pkg/agents/console/agent.go deleted file mode 100644 index b9ad907..0000000 --- a/pkg/agents/console/agent.go +++ /dev/null @@ -1,305 +0,0 @@ -package console - -import ( - "context" - "fmt" - "log/slog" - "os" - "path/filepath" - "strings" - - "github.com/docker/docker/api/types/container" - "github.com/docker/docker/api/types/mount" - "github.com/docker/docker/client" - - "gitea.stevedudenhoeffer.com/steve/answer/pkg/agents" - "gitea.stevedudenhoeffer.com/steve/answer/pkg/agents/shared" - gollm "gitea.stevedudenhoeffer.com/steve/go-llm" -) - -type Agent struct { - agents.Agent - - // Model is the chat completion model to use - Model gollm.ChatCompletion - - OnLoopComplete func(ctx context.Context, knowledge agents.Knowledge) error - - OnCommandStart func(ctx context.Context, command string) error - OnCommandDone func(ctx context.Context, command string, output string, err error) error - - OnDone func(ctx context.Context, knowledge agents.Knowledge) error - - ContextualInformation []string - - MaxCommands int -} - -type Response struct { - Knowledge agents.Knowledge - Directory string -} - -// Answer will give the model access to an ubuntu console with python and pip installed, and then ask the model to -// do what is necessary to answer the question. -func (a Agent) Answer(ctx context.Context, questions []string) (Response, error) { - var res Response - a.Agent = agents.NewAgent(a.Model, gollm.NewToolBox()).WithMaxCalls(200) - - if a.MaxCommands <= 0 { - a.MaxCommands = 10000 - } - - res.Knowledge = agents.Knowledge{ - OriginalQuestions: questions, - RemainingQuestions: questions, - } - - // create a temporary scratch directory - dir, err := os.MkdirTemp("", "console-") - if err != nil { - return res, err - } - - res.Directory = dir - - cl, err := client.NewClientWithOpts(client.FromEnv) - if err != nil { - return res, err - } - defer cl.Close() - - mounts := []mount.Mount{ - { - Type: mount.TypeBind, - Source: dir, - Target: "/home/user", - }, - } - - c, err := CreateContainer(ctx, cl, ContainerConfig{ - Config: &container.Config{ - Image: "ubuntu:latest", - Cmd: []string{"tail", "-f", "/dev/null"}, - Tty: true, - WorkingDir: "/home/user", - }, - HostConfig: &container.HostConfig{ - AutoRemove: true, - Mounts: mounts, - }, - Name: filepath.Base(dir), - }) - - if err != nil { - return res, fmt.Errorf("failed to create container: %w", err) - } - defer func() { - _ = c.Close(ctx) - }() - - slog.Info("starting container", "dir", dir, "container", fmt.Sprintf("%+v", c)) - err = c.Start(ctx) - if err != nil { - return res, fmt.Errorf("failed to start container: %w", err) - } - - // Run the model - - var history executions - var keepGoing = true - - opwd, epwd := c.Execute(ctx, "ls -al /home") - - fmt.Println(opwd) - slog.Info("pwd", "pwd", opwd, "epwd", epwd) - - tools := map[string]gollm.Function{ - "exit": gollm.NewFunction( - "exit", - "exit the container", - func(ctx *gollm.Context, args struct { - RemainingQuestions []string `description:"any remaining questions that remain unanswered"` - }) (any, error) { - keepGoing = false - return "exiting", nil - }), - - "write": gollm.NewFunction( - "write", - "write a file in the /root directory", - func(ctx *gollm.Context, args struct { - Filename string `description:"The name of the file to write"` - Content string `description:"The content of the file to write"` - }) (any, error) { - target, err := SafeJoinPath(dir, args.Filename) - if err != nil { - return "", err - } - - f, err := os.Create(target) - if err != nil { - return "", err - } - defer f.Close() - - _, err = f.WriteString(args.Content) - if err != nil { - return "", err - } - - return "wrote file", nil - }), - - "read": gollm.NewFunction( - "read", - "read a file in the /root directory", - func(ctx *gollm.Context, args struct { - Filename string `description:"The name of the file to read"` - }) (any, error) { - target, err := SafeJoinPath(dir, args.Filename) - if err != nil { - return "", err - } - - b, err := os.ReadFile(target) - if err != nil { - return "", err - } - - return string(b), nil - }), - - "execute": gollm.NewFunction( - "execute", - "execute a command in the container", - func(ctx *gollm.Context, args struct { - Command string `description:"The command to execute"` - }) (any, error) { - if len(history) >= a.MaxCommands { - return "too many commands", nil - } - - if a.OnCommandStart != nil { - err := a.OnCommandStart(ctx, args.Command) - if err != nil { - return "", err - } - } - - var res string - // if the command starts with sudo then we need to use the sudo function - if strings.HasPrefix(args.Command, "sudo ") { - res, err = c.Sudo(ctx, args.Command[5:]) - } else { - res, err = c.Execute(ctx, args.Command) - } - - if a.OnCommandDone != nil { - err = a.OnCommandDone(ctx, args.Command, res, nil) - if err != nil { - return "", err - } - } - - history = append(history, execution{ - Command: args.Command, - Output: res, - }) - - return res, nil - }), - - "sudo": gollm.NewFunction( - "sudo", - "execute a command in the container", - func(ctx *gollm.Context, args struct { - Command string `description:"The command to execute"` - }) (any, error) { - if len(history) >= a.MaxCommands { - return "too many commands", nil - } - - if a.OnCommandStart != nil { - err := a.OnCommandStart(ctx, args.Command) - if err != nil { - return "", err - } - } - - res, err := c.Sudo(ctx, args.Command) - - if a.OnCommandDone != nil { - err = a.OnCommandDone(ctx, args.Command, res, nil) - if err != nil { - return "", err - } - } - - if err != nil { - res = "error executing: " + err.Error() - } - - history = append(history, execution{ - Command: "sudo " + args.Command, - Output: res, - }) - - return res, nil - }), - } - - for i := 0; i < a.MaxCommands && len(history) < a.MaxCommands && keepGoing; i++ { - - systemPrompt := `You are now in a shell in a container of the ubuntu:latest image to answer a question asked by the user, it is very basic install of ubuntu, simple things (like python) are not preinstalled but can be installed via apt. You will be run multiple times and gain knowledge throughout the process.` - - if len(history) < a.MaxCommands { - systemPrompt += `You can run any command you like to get to the needed results.` - } - - systemPrompt += `Alternatively, you can use the tool "write" to write a file in the home directory, and also the tool "read" to read a file in the home directory. -When you are done, please use "exit" to exit the container. -Respond with any number of commands to answer the question, they will be executed in order.` - - var toolbox []gollm.Function - - // add unrestricted tools - toolbox = append(toolbox, tools["exit"], tools["write"], tools["read"]) - - if len(history) < a.MaxCommands { - toolbox = append(toolbox, tools["execute"], tools["sudo"]) - } - - kw := shared.KnowledgeWorker{ - Model: a.Model, - ToolBox: gollm.NewToolBox(toolbox...), - ContextualInformation: a.ContextualInformation, - OnNewFunction: func(ctx context.Context, funcName string, args string) (any, error) { - slog.Info("new function called", "function name", funcName, "args", args) - return nil, nil - }, - } - - slog.Info("answering question", "question", questions[0]) - r, err := kw.Answer(ctx, &res.Knowledge, systemPrompt, "", "", history.ToGeneralButLastMessageHistory(), func(res gollm.ToolCallResponse) { - }) - - if err != nil { - return res, fmt.Errorf("error answering question: %w", err) - } - - if len(r.Knowledge) > 0 { - slog.Info("answered question and learned", "knowledge", r.Knowledge) - } else { - slog.Info("answered question and learned nothing") - } - - res.Knowledge, err = a.KnowledgeIntegrate(ctx, res.Knowledge, r) - if err != nil { - return res, fmt.Errorf("error integrating knowledge: %w", err) - } - - slog.Info("knowledge integrated", "question", questions[0], "knowledge", res.Knowledge) - } - return res, nil -} diff --git a/pkg/agents/console/execution.go b/pkg/agents/console/execution.go index 4345abf..4d41eaf 100644 --- a/pkg/agents/console/execution.go +++ b/pkg/agents/console/execution.go @@ -6,7 +6,7 @@ import ( gollm "gitea.stevedudenhoeffer.com/steve/go-llm" ) -type execution struct { +type Execution struct { Command string Output string WhatILearned []string @@ -16,7 +16,7 @@ type execution struct { const kMaxLenCommandSummary = 200 const kMaxLenCommandOutputSummary = 200 -func (e execution) ToGeneralMessageHistory() gollm.Message { +func (e Execution) ToGeneralMessageHistory() gollm.Message { if len(e.Command) > kMaxLenCommandSummary { e.Command = e.Command[:kMaxLenCommandSummary] + "... (truncated)" } @@ -33,7 +33,7 @@ func (e execution) ToGeneralMessageHistory() gollm.Message { } } -func (e execution) ToDetailedMessageHistory() gollm.Message { +func (e Execution) ToDetailedMessageHistory() gollm.Message { prompt := "$ " if strings.HasPrefix(e.Command, "sudo ") { prompt = "# " @@ -60,9 +60,9 @@ func (e execution) ToDetailedMessageHistory() gollm.Message { } } -type executions []execution +type Executions []Execution -func (e executions) ToGeneralMessageHistory() []gollm.Message { +func (e Executions) ToGeneralMessageHistory() []gollm.Message { var messages []gollm.Message for _, v := range e { @@ -72,7 +72,7 @@ func (e executions) ToGeneralMessageHistory() []gollm.Message { return messages } -func (e executions) ToGeneralButLastMessageHistory() []gollm.Message { +func (e Executions) ToGeneralButLastMessageHistory() []gollm.Message { var messages []gollm.Message for i, v := range e { @@ -86,7 +86,7 @@ func (e executions) ToGeneralButLastMessageHistory() []gollm.Message { return messages } -func (e executions) ToDetailedMessageHistory() []gollm.Message { +func (e Executions) ToDetailedMessageHistory() []gollm.Message { var messages []gollm.Message for _, v := range e {