cmd/mimi: try having mimi run Python

Signed-off-by: Xe Iaso <me@xeiaso.net>

6 files changed, 209 insertions, 0 deletions

diff --git a/cmd/mimi/modules/discord/jufra/jufra.go b/cmd/mimi/modules/discord/jufra/jufra.go
index 43caeb3..6772a07 100644
--- a/cmd/mimi/modules/discord/jufra/jufra.go
+++ b/cmd/mimi/modules/discord/jufra/jufra.go
@@ -225,6 +225,7 @@ func (m *Module) messageCreate(s *discordgo.Session, mc *discordgo.MessageCreate
 		Options: map[string]any{
 			"num_ctx": 131072,
 		},
+		Tools: m.getTools(),
 	}
 
 	resp, err := m.ollama.Chat(context.Background(), cr)
@@ -236,6 +237,37 @@ func (m *Module) messageCreate(s *discordgo.Session, mc *discordgo.MessageCreate
 
 	conv = append(conv, resp.Message)
 
+	if len(resp.Message.ToolCalls) != 0 {
+		for _, tc := range resp.Message.ToolCalls {
+			if tc.Name == "run_python_code" {
+				msg, err := m.runPythonCode(context.Background(), tc)
+				if err != nil {
+					slog.Error("error running python code", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
+					s.ChannelMessageSend(mc.ChannelID, "error running python code")
+					return
+				}
+
+				conv = append(conv, *msg)
+
+				resp, err = m.ollama.Chat(context.Background(), &ollama.CompleteRequest{
+					Model:    *mimiModel,
+					Messages: conv,
+					Options: map[string]any{
+						"num_ctx": 131072,
+					},
+					Tools: m.getTools(),
+				})
+				if err != nil {
+					slog.Error("error chatting", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
+					s.ChannelMessageSend(mc.ChannelID, "error chatting")
+					return
+				}
+
+				conv = append(conv, resp.Message)
+			}
+		}
+	}
+
 	if !*disableLlamaguard {
 		lgResp, err := m.llamaGuardCheck(context.Background(), "assistant", conv)
 		if err != nil {
diff --git a/cmd/mimi/modules/discord/jufra/tools.go b/cmd/mimi/modules/discord/jufra/tools.go
index 9fdcaca..4c82c68 100644
--- a/cmd/mimi/modules/discord/jufra/tools.go
+++ b/cmd/mimi/modules/discord/jufra/tools.go
@@ -1 +1,77 @@
 package jufra
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"os"
+
+	"within.website/x/llm/codeinterpreter/python"
+	"within.website/x/web/ollama"
+)
+
+var normalTools = []ollama.Function{
+	{
+		Name:        "run_python_code",
+		Description: "Run the given Python code in a sandboxed environment",
+		Parameters: ollama.Param{
+			Type: "object",
+			Properties: ollama.Properties{
+				"code": {
+					Type:        "string",
+					Description: "The Python code to run",
+				},
+			},
+			Required: []string{"code"},
+		},
+	},
+}
+
+type pythonCodeArgs struct {
+	Code string `json:"code"`
+}
+
+func (pca *pythonCodeArgs) Valid() error {
+	if pca.Code == "" {
+		return errors.New("missing code parameter")
+	}
+
+	return nil
+}
+
+func (m *Module) runPythonCode(ctx context.Context, tc ollama.ToolCall) (*ollama.Message, error) {
+	var args pythonCodeArgs
+	if err := json.Unmarshal(tc.Arguments, &args); err != nil {
+		return nil, err
+	}
+
+	tmpdir, err := os.MkdirTemp("", "mimi-python-*")
+	if err != nil {
+		return nil, err
+	}
+
+	defer os.RemoveAll(tmpdir)
+
+	res, err := python.Run(ctx, tmpdir, args.Code)
+	if err != nil {
+		return nil, nil
+	}
+
+	return &ollama.Message{
+		Role:    "tool",
+		Content: jsonString(res),
+	}, nil
+}
+
+func (m *Module) getTools() []ollama.Tool {
+	var result []ollama.Tool
+
+	for _, tool := range normalTools {
+		result = append(result, ollama.Tool{
+			Type:     "function",
+			Function: tool,
+		})
+	}
+
+	return result
+}
diff --git a/llm/codeinterpreter/python/output/.gitignore b/llm/codeinterpreter/python/output/.gitignore
new file mode 100644
index 0000000..c96a04f
--- /dev/null
+++ b/llm/codeinterpreter/python/output/.gitignore
@@ -0,0 +1,2 @@
+*
+!.gitignore
\ No newline at end of file
diff --git a/llm/codeinterpreter/python/python.go b/llm/codeinterpreter/python/python.go
new file mode 100644
index 0000000..06457b1
--- /dev/null
+++ b/llm/codeinterpreter/python/python.go
@@ -0,0 +1,77 @@
+package python
+
+import (
+	"bytes"
+	"context"
+	_ "embed"
+	"fmt"
+	"os"
+
+	"github.com/tetratelabs/wazero"
+	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
+)
+
+var (
+	//go:embed python.wasm
+	Binary []byte
+
+	r    wazero.Runtime
+	code wazero.CompiledModule
+)
+
+func init() {
+	ctx := context.Background()
+	r = wazero.NewRuntime(ctx)
+
+	wasi_snapshot_preview1.MustInstantiate(ctx, r)
+
+	var err error
+	code, err = r.CompileModule(ctx, Binary)
+	if err != nil {
+		panic(err)
+	}
+}
+
+type Result struct {
+	Stdout string
+	Stderr string
+}
+
+func Run(ctx context.Context, tmpDir, userCode string) (*Result, error) {
+	fout := &bytes.Buffer{}
+	ferr := &bytes.Buffer{}
+	fin := &bytes.Buffer{}
+
+	os.WriteFile(tmpDir+"/main.py", []byte(userCode), 0644)
+
+	fsConfig := wazero.NewFSConfig().
+		WithFSMount(os.DirFS(tmpDir), "/")
+
+	config := wazero.NewModuleConfig().
+		// stdio
+		WithStdout(fout).
+		WithStderr(ferr).
+		WithStdin(fin).
+		// argv
+		WithArgs("python", "/main.py").
+		WithName("python").
+		// fs / system
+		WithFSConfig(fsConfig).
+		WithSysNanosleep().
+		WithSysNanotime().
+		WithSysWalltime()
+
+	mod, err := r.InstantiateModule(ctx, code, config)
+	if err != nil {
+		fmt.Println(fout.String())
+		fmt.Println(ferr.String())
+		return nil, err
+	}
+
+	defer mod.Close(ctx)
+
+	return &Result{
+		Stdout: fout.String(),
+		Stderr: ferr.String(),
+	}, nil
+}
diff --git a/llm/codeinterpreter/python/python.wasm b/llm/codeinterpreter/python/python.wasm
new file mode 100644
index 0000000..159a20b
--- /dev/null
+++ b/llm/codeinterpreter/python/python.wasm
diff --git a/llm/codeinterpreter/python/python_test.go b/llm/codeinterpreter/python/python_test.go
new file mode 100644
index 0000000..7b22f65
--- /dev/null
+++ b/llm/codeinterpreter/python/python_test.go
@@ -0,0 +1,22 @@
+package python
+
+import (
+	"context"
+	"testing"
+)
+
+func TestRun(t *testing.T) {
+	var code = `import sys
+
+print(f"Python {sys.version} running in {sys.platform}/wazero.")`
+
+	dir := t.TempDir()
+
+	res, err := Run(context.Background(), dir, code)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Logf("stdout: %s", res.Stdout)
+	t.Logf("stderr: %s", res.Stderr)
+}