Moved the Knowledge struct and related types to the shared package, updating all references across the codebase. This improves modularity and enables better reuse of the Knowledge type across different components.
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 shared.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 shared.Knowledge) error
|
|
|
|
ContextualInformation []string
|
|
|
|
MaxCommands int
|
|
}
|
|
|
|
type Response struct {
|
|
Knowledge shared.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 = shared.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
|
|
}
|