Steve Dudenhoeffer 1c3ea7d1f1 Refactor Knowledge struct into shared package
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.
2025-05-03 22:09:02 -04:00

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
}