303 lines
7.7 KiB
Go
303 lines
7.7 KiB
Go
package console_new
|
|
|
|
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 {
|
|
// 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
|
|
|
|
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 = agents.KnowledgeIntegrate(ctx, a.Model, 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
|
|
}
|