diff --git a/llm/cli.py b/llm/cli.py index 2e11e2c8..57f73828 100644 --- a/llm/cli.py +++ b/llm/cli.py @@ -673,8 +673,12 @@ def read_prompt(): template_obj = load_template(template) except LoadTemplateError as ex: raise click.ClickException(str(ex)) - extract = template_obj.extract - extract_last = template_obj.extract_last + extract = template_obj.extract if template_obj.extract is not None else extract + extract_last = ( + template_obj.extract_last + if template_obj.extract_last is not None + else extract_last + ) # Combine with template fragments/system_fragments if template_obj.fragments: fragments = [*template_obj.fragments, *fragments] diff --git a/tests/test_templates.py b/tests/test_templates.py index dcdabd9b..4f6d6250 100644 --- a/tests/test_templates.py +++ b/tests/test_templates.py @@ -546,3 +546,42 @@ def after(): finally: after() pm.unregister(name="greetings-plugin") + + +@mock.patch.dict(os.environ, {"OPENAI_API_KEY": "X"}) +def test_extract_last_flag_with_templates(templates_path, mocked_openai_chat): + """Test that --xl flag works with templates - issue #1296""" + # Create a template without extract_last set + (templates_path / "code-template.yaml").write_text( + "prompt: 'Write Python code for $input'", "utf-8" + ) + runner = CliRunner() + # Test with --xl flag + result = runner.invoke( + cli, + ["--no-stream", "-t", "code-template", "hello world", "--xl"], + catch_exceptions=False, + ) + assert result.exit_code == 0 + # The output should be extracted code, not the full response + # Check that mocked response extraction was applied + assert "```" not in result.output.strip() + + +@mock.patch.dict(os.environ, {"OPENAI_API_KEY": "X"}) +def test_extract_last_in_template_with_cli_flag(templates_path, mocked_openai_chat): + """Test that CLI --xl flag works even when template has extract_last: false""" + # Create a template with extract_last explicitly set to false + (templates_path / "no-extract-template.yaml").write_text( + "prompt: 'Write Python code for $input'\nextract_last: false", "utf-8" + ) + runner = CliRunner() + # Test with --xl flag - should override template setting + result = runner.invoke( + cli, + ["--no-stream", "-t", "no-extract-template", "hello world", "--xl"], + catch_exceptions=False, + ) + assert result.exit_code == 0 + # The output should be extracted code since CLI flag should take precedence + assert "```" not in result.output.strip()