Merge pull request #57 from semanser/fix-tool-placeholder-issue

Fix tool placeholder issue

Author: semanser

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 {

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
 	}

backend/providers/types.go

@@ -2,15 +2,28 @@ 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 {
 	Input string
 	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
+}