answer/pkg/agents/console.go

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"
)

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
}