From 23e68d7be39f148bf32cf2f4f1220d671ef34bf3 Mon Sep 17 00:00:00 2001
From: Andriy Semenets <semanser@gmail.com>
Date: Fri, 5 Apr 2024 12:32:27 +0200
Subject: [PATCH 1/2] Fix an issue with empty ToolPlaceholder

---
 backend/providers/openai.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/backend/providers/openai.go b/backend/providers/openai.go
index 88e882e..f73ce65 100644
--- a/backend/providers/openai.go
+++ b/backend/providers/openai.go
@@ -57,7 +57,13 @@ func (p OpenAIProvider) DockerImageName(task string) (string, error) {
 func (p OpenAIProvider) NextTask(args NextTaskOptions) *database.Task {
 	log.Println("Getting next task")
 
-	prompt, err := templates.Render(assets.PromptTemplates, "prompts/agent.tmpl", args)
+	promptArgs := map[string]interface{}{
+		"DockerImage":     args.DockerImage,
+		"ToolPlaceholder": getToolPlaceholder(),
+		"Tasks":           args.Tasks,
+	}
+
+	prompt, err := templates.Render(assets.PromptTemplates, "prompts/agent.tmpl", promptArgs)
 
 	// TODO In case of lots of tasks, we should try to get a summary using gpt-3.5
 	if len(prompt) > 30000 {

From 0c322fce1666d53b45913583fa7f5883ef4cf7ea Mon Sep 17 00:00:00 2001
From: Andriy Semenets <semanser@gmail.com>
Date: Fri, 5 Apr 2024 12:49:17 +0200
Subject: [PATCH 2/2] Fix an issue with printing messages to the terminal

---
 backend/providers/providers.go | 23 +++++++++++++++++++++--
 backend/providers/types.go     | 29 +++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/backend/providers/providers.go b/backend/providers/providers.go
index cdc5a92..6ce9f7e 100644
--- a/backend/providers/providers.go
+++ b/backend/providers/providers.go
@@ -249,7 +249,26 @@ func toolToTask(choices []*llms.ContentChoice) (*database.Task, error) {
 	}
 
 	// We use AskArgs to extract the message
-	params, err := extractToolArgs(tool.FunctionCall.Arguments, &AskArgs{})
+	var toolType Messanger
+
+	switch tool.FunctionCall.Name {
+	case "input":
+		toolType = &InputArgs{}
+	case "terminal":
+		toolType = &TerminalArgs{}
+	case "browser":
+		toolType = &BrowserArgs{}
+	case "code":
+		toolType = &CodeArgs{}
+	case "ask":
+		toolType = &AskArgs{}
+	case "done":
+		toolType = &DoneArgs{}
+	default:
+		return nil, fmt.Errorf("unknown tool name: %s", tool.FunctionCall.Name)
+	}
+
+	params, err := extractToolArgs(tool.FunctionCall.Arguments, &toolType)
 	if err != nil {
 		return nil, fmt.Errorf("failed to extract args: %v", err)
 	}
@@ -260,7 +279,7 @@ func toolToTask(choices []*llms.ContentChoice) (*database.Task, error) {
 	task.Args = database.StringToNullString(string(args))
 
 	// Sometimes the model returns an empty string for the message
-	msg := string(params.Message)
+	msg := string((*params).GetMessage())
 	if msg == "" {
 		msg = tool.FunctionCall.Arguments
 	}
diff --git a/backend/providers/types.go b/backend/providers/types.go
index 3bc27fc..0b73705 100644
--- a/backend/providers/types.go
+++ b/backend/providers/types.go
@@ -2,8 +2,17 @@ package providers
 
 type Message string
 
+type Messanger interface {
+	GetMessage() Message
+}
+
 type InputArgs struct {
 	Query string
+	Message
+}
+
+func (i *InputArgs) GetMessage() Message {
+	return i.Message
 }
 
 type TerminalArgs struct {
@@ -11,6 +20,10 @@ type TerminalArgs struct {
 	Message
 }
 
+func (t *TerminalArgs) GetMessage() Message {
+	return t.Message
+}
+
 type BrowserAction string
 
 const (
@@ -24,6 +37,10 @@ type BrowserArgs struct {
 	Message
 }
 
+func (b *BrowserArgs) GetMessage() Message {
+	return b.Message
+}
+
 type CodeAction string
 
 const (
@@ -38,10 +55,22 @@ type CodeArgs struct {
 	Message
 }
 
+func (c *CodeArgs) GetMessage() Message {
+	return c.Message
+}
+
 type AskArgs struct {
 	Message
 }
 
+func (a *AskArgs) GetMessage() Message {
+	return a.Message
+}
+
 type DoneArgs struct {
 	Message
 }
+
+func (d *DoneArgs) GetMessage() Message {
+	return d.Message
+}