398 lines
11 KiB
Go
398 lines
11 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"
|
|
|
|
gollm "gitea.stevedudenhoeffer.com/steve/go-llm"
|
|
|
|
"gitea.stevedudenhoeffer.com/steve/answer/pkg/agents"
|
|
"gitea.stevedudenhoeffer.com/steve/answer/pkg/agents/shared"
|
|
)
|
|
|
|
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
|
|
DataDir string
|
|
OutputDir 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 = 20 // Default to 20 commands as per requirements
|
|
}
|
|
|
|
res.Knowledge = agents.Knowledge{
|
|
OriginalQuestions: questions,
|
|
RemainingQuestions: questions,
|
|
}
|
|
|
|
// create temporary directories for data and output
|
|
dataDir, err := os.MkdirTemp("", "console-data-")
|
|
if err != nil {
|
|
return res, err
|
|
}
|
|
|
|
outputDir, err := os.MkdirTemp("", "console-output-")
|
|
if err != nil {
|
|
os.RemoveAll(dataDir)
|
|
return res, err
|
|
}
|
|
|
|
res.DataDir = dataDir
|
|
res.OutputDir = outputDir
|
|
|
|
cl, err := client.NewClientWithOpts(client.FromEnv)
|
|
if err != nil {
|
|
os.RemoveAll(dataDir)
|
|
os.RemoveAll(outputDir)
|
|
return res, err
|
|
}
|
|
defer cl.Close()
|
|
|
|
mounts := []mount.Mount{
|
|
{
|
|
Type: mount.TypeBind,
|
|
Source: dataDir,
|
|
Target: "/data",
|
|
},
|
|
{
|
|
Type: mount.TypeBind,
|
|
Source: outputDir,
|
|
Target: "/output",
|
|
},
|
|
}
|
|
|
|
c, err := CreateContainer(ctx, cl, ContainerConfig{
|
|
Config: &container.Config{
|
|
Image: "ubuntu:latest",
|
|
Cmd: []string{"tail", "-f", "/dev/null"},
|
|
Tty: true,
|
|
WorkingDir: "/data",
|
|
},
|
|
HostConfig: &container.HostConfig{
|
|
AutoRemove: true,
|
|
Mounts: mounts,
|
|
},
|
|
Name: filepath.Base(dataDir),
|
|
})
|
|
|
|
if err != nil {
|
|
os.RemoveAll(dataDir)
|
|
os.RemoveAll(outputDir)
|
|
return res, fmt.Errorf("failed to create container: %w", err)
|
|
}
|
|
defer func() {
|
|
_ = c.Close(ctx)
|
|
}()
|
|
|
|
slog.Info("starting container", "dataDir", dataDir, "outputDir", outputDir, "container", fmt.Sprintf("%+v", c))
|
|
err = c.Start(ctx)
|
|
if err != nil {
|
|
os.RemoveAll(dataDir)
|
|
os.RemoveAll(outputDir)
|
|
return res, fmt.Errorf("failed to start container: %w", err)
|
|
}
|
|
|
|
// Run the model
|
|
|
|
var history executions
|
|
var keepGoing = true
|
|
|
|
// Initial setup - install basic tools
|
|
setupCmd, setupErr := c.Sudo(ctx, "apt-get update && apt-get install -y curl wget git python3 python3-pip")
|
|
if setupErr == nil {
|
|
history = append(history, execution{
|
|
Command: "sudo apt-get update && apt-get install -y curl wget git python3 python3-pip",
|
|
Output: setupCmd,
|
|
WhatILearned: []string{"Basic tools installed: curl, wget, git, python3, pip"},
|
|
})
|
|
}
|
|
|
|
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_data": gollm.NewFunction(
|
|
"write_data",
|
|
"write a file in the /data 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(dataDir, args.Filename)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// Ensure directory exists
|
|
dir := filepath.Dir(target)
|
|
if err := os.MkdirAll(dir, 0755); 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 to /data/" + args.Filename, nil
|
|
}),
|
|
|
|
"write_output": gollm.NewFunction(
|
|
"write_output",
|
|
"write a file in the /output directory (files here will be available after the agent completes)",
|
|
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(outputDir, args.Filename)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// Ensure directory exists
|
|
dir := filepath.Dir(target)
|
|
if err := os.MkdirAll(dir, 0755); 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 to /output/" + args.Filename, nil
|
|
}),
|
|
|
|
"read_data": gollm.NewFunction(
|
|
"read_data",
|
|
"read a file from the /data directory",
|
|
func(ctx *gollm.Context, args struct {
|
|
Filename string `description:"The name of the file to read"`
|
|
}) (any, error) {
|
|
target, err := SafeJoinPath(dataDir, args.Filename)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
b, err := os.ReadFile(target)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return string(b), nil
|
|
}),
|
|
|
|
"read_output": gollm.NewFunction(
|
|
"read_output",
|
|
"read a file from the /output directory",
|
|
func(ctx *gollm.Context, args struct {
|
|
Filename string `description:"The name of the file to read"`
|
|
}) (any, error) {
|
|
target, err := SafeJoinPath(outputDir, 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"`
|
|
Learn []string `description:"What you learned from this command (optional)"`
|
|
ToLearn []string `description:"What you still need to learn (optional)"`
|
|
}) (any, error) {
|
|
if len(history) >= a.MaxCommands {
|
|
return "Command limit reached. You've used all " + fmt.Sprintf("%d", a.MaxCommands) + " allowed 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,
|
|
WhatILearned: args.Learn,
|
|
WhatIStillNeedToLearn: args.ToLearn,
|
|
})
|
|
|
|
return res, nil
|
|
}),
|
|
|
|
"summarize_knowledge": gollm.NewFunction(
|
|
"summarize_knowledge",
|
|
"summarize what you've learned so far",
|
|
func(ctx *gollm.Context, args struct{}) (any, error) {
|
|
var learned []string
|
|
var toLearn []string
|
|
|
|
for _, exec := range history {
|
|
learned = append(learned, exec.WhatILearned...)
|
|
toLearn = append(toLearn, exec.WhatIStillNeedToLearn...)
|
|
}
|
|
|
|
summary := "Knowledge Summary:\n"
|
|
if len(learned) > 0 {
|
|
summary += "What I've learned:\n- " + strings.Join(learned, "\n- ") + "\n\n"
|
|
} else {
|
|
summary += "I haven't learned anything specific yet.\n\n"
|
|
}
|
|
|
|
if len(toLearn) > 0 {
|
|
summary += "What I still need to learn:\n- " + strings.Join(toLearn, "\n- ")
|
|
} else {
|
|
summary += "I don't have any specific learning goals at the moment."
|
|
}
|
|
|
|
return summary, 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.
|
|
You have full control over a bash shell inside this Docker container.
|
|
|
|
Important directories:
|
|
- /data: A temporary directory with the lifespan of your processing. Use this for working files.
|
|
- /output: Files placed here will be returned to the caller after you're done. Use this for final results.
|
|
|
|
You can execute up to ` + fmt.Sprintf("%d", a.MaxCommands) + ` commands total. You're currently on command ` + fmt.Sprintf("%d", len(history)+1) + ` of ` + fmt.Sprintf("%d", a.MaxCommands) + `.
|
|
|
|
For each command, you should:
|
|
1. Think about what you need to learn
|
|
2. Execute the command using the "execute" function
|
|
3. Analyze the output and record what you learned
|
|
4. Plan your next command based on this knowledge
|
|
|
|
You can write files directly to /data or /output using the write_data and write_output functions.
|
|
When you are done, use "exit" to exit the container.`
|
|
|
|
var toolbox []gollm.Function
|
|
|
|
// Add all tools
|
|
toolbox = append(toolbox,
|
|
tools["exit"],
|
|
tools["write_data"],
|
|
tools["write_output"],
|
|
tools["read_data"],
|
|
tools["read_output"],
|
|
tools["summarize_knowledge"],
|
|
)
|
|
|
|
// Only add execute if we haven't reached the command limit
|
|
if len(history) < a.MaxCommands {
|
|
toolbox = append(toolbox, tools["execute"])
|
|
}
|
|
|
|
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
|
|
}
|