diff --git a/llm_cmd.py b/llm_cmd.py
index 33b9d33..348d33a 100644
--- a/llm_cmd.py
+++ b/llm_cmd.py
@@ -1,5 +1,6 @@
import click
import llm
+import re
import subprocess
from prompt_toolkit import PromptSession
from prompt_toolkit.lexers import PygmentsLexer
@@ -14,6 +15,12 @@
You return only: git reset --soft HEAD~1
""".strip()
+THINKING_RE = re.compile(r".*?", re.DOTALL)
+
+def strip_thinking(text):
+ # Thinking models may emit ... blocks alongside the command
+ return THINKING_RE.sub("", text).strip()
+
@llm.hookimpl
def register_commands(cli):
@cli.command()
@@ -30,7 +37,7 @@ def cmd(args, model, system, key):
if model_obj.needs_key:
model_obj.key = llm.get_key(key, model_obj.needs_key, model_obj.key_env_var)
result = model_obj.prompt(prompt, system=system or SYSTEM_PROMPT)
- interactive_exec(str(result))
+ interactive_exec(strip_thinking(str(result)))
def interactive_exec(command):
session = PromptSession(lexer=PygmentsLexer(BashLexer))
diff --git a/tests/test_cmd.py b/tests/test_cmd.py
index 02b5db8..d5a2ab3 100644
--- a/tests/test_cmd.py
+++ b/tests/test_cmd.py
@@ -1,6 +1,18 @@
-from llm.plugins import pm
+from llm.plugins import load_plugins, pm
+from llm_cmd import strip_thinking
def test_plugin_is_installed():
+ load_plugins()
names = [mod.__name__ for mod in pm.get_plugins()]
assert "llm_cmd" in names
+
+
+def test_strip_thinking():
+ # Response shape from https://github.com/simonw/llm-cmd/issues/27
+ text = "\nOkay, the user wants to show their bash history.\n\n\n history\n"
+ assert strip_thinking(text) == "history"
+
+
+def test_strip_thinking_no_block():
+ assert strip_thinking("ls -la") == "ls -la"