From f832183487b1c5c280ccba0e2478121093a18804 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Wed, 18 Dec 2024 17:12:58 +0100
Subject: [PATCH] fix: properly handle piped input in prompts (#354)

Previously, piped input was treated as a separate message instead of being included in the user prompt. This change:
- Adds inject_stdin function to handle piped input
- Ensures piped input is appended to first prompt or creates new message if none exists
- Removes redundant system message creation
---
 gptme/cli.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/gptme/cli.py b/gptme/cli.py
index d30ed5b9..09ac063c 100644
--- a/gptme/cli.py
+++ b/gptme/cli.py
@@ -238,10 +238,22 @@ def main(
     # TODO: referenced file paths in multiprompts should be read when run, not when parsed
     prompt_msgs = [Message("user", p) for p in prompts]
 
+    def inject_stdin(prompt_msgs, piped_input: str | None) -> list[Message]:
+        # if piped input, append it to first prompt, or create a new prompt if none exists
+        if not piped_input:
+            return prompt_msgs
+        stdin_msg = Message("user", f"```stdin\n{piped_input}\n```")
+        if not prompt_msgs:
+            prompt_msgs.append(stdin_msg)
+        else:
+            prompt_msgs[0] = prompt_msgs[0].replace(
+                content=f"{prompt_msgs[0].content}\n\n{stdin_msg.content}"
+            )
+        return prompt_msgs
+
     if resume:
         logdir = get_logdir_resume()
-        if piped_input:
-            prompt_msgs.append(Message("system", f"```stdin\n{piped_input}\n```"))
+        prompt_msgs = inject_stdin(prompt_msgs, piped_input)
     # don't run pick in tests/non-interactive mode, or if the user specifies a name
     elif (
         interactive
@@ -253,8 +265,7 @@ def main(
         logdir = pick_log()
     else:
         logdir = get_logdir(name)
-        if piped_input:
-            initial_msgs.append(Message("system", f"```stdin\n{piped_input}\n```"))
+        prompt_msgs = inject_stdin(prompt_msgs, piped_input)
 
     if workspace == "@log":
         workspace_path: Path | None = logdir / "workspace"