Allow templating {{input}} in prompt files (#46)

Author: sgoedecke

cmd/run/run.go

@@ -260,6 +260,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 
 			initialPrompt := ""
 			singleShot := false
+			pipedContent := ""
 
 			if len(args) > 1 {
 				initialPrompt = strings.Join(args[1:], " ")
@@ -269,8 +270,11 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 			if isPipe(os.Stdin) {
 				promptFromPipe, _ := io.ReadAll(os.Stdin)
 				if len(promptFromPipe) > 0 {
-					initialPrompt = initialPrompt + "\n" + string(promptFromPipe)
-					singleShot = true
+					pipedContent = strings.TrimSpace(string(promptFromPipe))
+					if pf == nil {
+						initialPrompt = initialPrompt + "\n" + pipedContent
+						singleShot = true
+					}
 				}
 			}
 
@@ -283,22 +287,28 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 				systemPrompt: systemPrompt,
 			}
 
-			// preload conversation & parameters from YAML
+			// If a prompt file is passed, load the messages from the file, templating {{input}} from stdin
 			if pf != nil {
 				for _, m := range pf.Messages {
+					content := m.Content
+					if pipedContent != "" && strings.ToLower(m.Role) == "user" {
+						content = strings.ReplaceAll(content, "{{input}}", pipedContent)
+					}
 					switch strings.ToLower(m.Role) {
 					case "system":
 						if conversation.systemPrompt == "" {
-							conversation.systemPrompt = m.Content
+							conversation.systemPrompt = content
 						} else {
-							conversation.AddMessage(azuremodels.ChatMessageRoleSystem, m.Content)
+							conversation.AddMessage(azuremodels.ChatMessageRoleSystem, content)
 						}
 					case "user":
-						conversation.AddMessage(azuremodels.ChatMessageRoleUser, m.Content)
+						conversation.AddMessage(azuremodels.ChatMessageRoleUser, content)
 					case "assistant":
-						conversation.AddMessage(azuremodels.ChatMessageRoleAssistant, m.Content)
+						conversation.AddMessage(azuremodels.ChatMessageRoleAssistant, content)
 					}
 				}
+
+				initialPrompt = ""
 			}
 
 			mp := ModelParameters{}
@@ -320,7 +330,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 					initialPrompt = ""
 				}
 
-				if prompt == "" {
+				if prompt == "" && pf == nil {
 					fmt.Printf(">>> ")
 					reader := bufio.NewReader(os.Stdin)
 					prompt, err = reader.ReadString('\n')
@@ -331,7 +341,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 
 				prompt = strings.TrimSpace(prompt)
 
-				if prompt == "" {
+				if prompt == "" && pf == nil {
 					continue
 				}
 
@@ -419,7 +429,7 @@ func NewRunCommand(cfg *command.Config) *cobra.Command {
 
 				conversation.AddMessage(azuremodels.ChatMessageRoleAssistant, messageBuilder.String())
 
-				if singleShot {
+				if singleShot || pf != nil {
 					break
 				}
 			}

cmd/run/run_test.go

@@ -134,7 +134,6 @@ messages:
 		runCmd.SetArgs([]string{
 			"--file", tmp.Name(),
 			azuremodels.FormatIdentifier("openai", "test-model"),
-			"foo?",
 		})
 
 		_, err = runCmd.ExecuteC()
@@ -143,11 +142,85 @@ messages:
 		require.Equal(t, 3, len(capturedReq.Messages))
 		require.Equal(t, "You are a text summarizer.", *capturedReq.Messages[0].Content)
 		require.Equal(t, "Hello there!", *capturedReq.Messages[1].Content)
-		require.Equal(t, "foo?", *capturedReq.Messages[2].Content)
 
 		require.NotNil(t, capturedReq.Temperature)
 		require.Equal(t, 0.5, *capturedReq.Temperature)
 
 		require.Contains(t, out.String(), reply) // response streamed to output
 	})
+
+	t.Run("--file with {{input}} placeholder is substituted with stdin", func(t *testing.T) {
+		const yamlBody = `
+name: Summarizer
+description: Summarizes input text
+model: openai/test-model
+messages:
+  - role: system
+    content: You are a text summarizer.
+  - role: user
+    content: "{{input}}"
+`
+
+		tmp, err := os.CreateTemp(t.TempDir(), "*.prompt.yml")
+		require.NoError(t, err)
+		_, err = tmp.WriteString(yamlBody)
+		require.NoError(t, err)
+		require.NoError(t, tmp.Close())
+
+		client := azuremodels.NewMockClient()
+		modelSummary := &azuremodels.ModelSummary{
+			Name:      "test-model",
+			Publisher: "openai",
+			Task:      "chat-completion",
+		}
+		client.MockListModels = func(ctx context.Context) ([]*azuremodels.ModelSummary, error) {
+			return []*azuremodels.ModelSummary{modelSummary}, nil
+		}
+
+		var capturedReq azuremodels.ChatCompletionOptions
+		reply := "Summary - bar"
+		chatCompletion := azuremodels.ChatCompletion{
+			Choices: []azuremodels.ChatChoice{{
+				Message: &azuremodels.ChatChoiceMessage{
+					Content: util.Ptr(reply),
+					Role:    util.Ptr(string(azuremodels.ChatMessageRoleAssistant)),
+				},
+			}},
+		}
+		client.MockGetChatCompletionStream = func(ctx context.Context, opt azuremodels.ChatCompletionOptions) (*azuremodels.ChatCompletionResponse, error) {
+			capturedReq = opt
+			return &azuremodels.ChatCompletionResponse{
+				Reader: sse.NewMockEventReader([]azuremodels.ChatCompletion{chatCompletion}),
+			}, nil
+		}
+
+		// create a pipe to fake stdin so that isPipe(os.Stdin)==true
+		r, w, err := os.Pipe()
+		require.NoError(t, err)
+		oldStdin := os.Stdin
+		os.Stdin = r
+		defer func() { os.Stdin = oldStdin }()
+		piped := "Hello there!"
+		go func() {
+			_, _ = w.Write([]byte(piped))
+			_ = w.Close()
+		}()
+
+		out := new(bytes.Buffer)
+		cfg := command.NewConfig(out, out, client, true, 100)
+		runCmd := NewRunCommand(cfg)
+		runCmd.SetArgs([]string{
+			"--file", tmp.Name(),
+			azuremodels.FormatIdentifier("openai", "test-model"),
+		})
+
+		_, err = runCmd.ExecuteC()
+		require.NoError(t, err)
+
+		require.Len(t, capturedReq.Messages, 3)
+		require.Equal(t, "You are a text summarizer.", *capturedReq.Messages[0].Content)
+		require.Equal(t, piped, *capturedReq.Messages[1].Content) // {{input}} -> "Hello there!"
+
+		require.Contains(t, out.String(), reply)
+	})
 }