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"