llm prompt -x/--extract option, closes #681

docs/help.md

@@ -105,6 +105,11 @@ Usage: llm prompt [OPTIONS] [PROMPT]
       # With an explicit mimetype:
       cat image | llm 'describe image' --at - image/jpeg
 
+  The -x/--extract option returns just the content of the first ``` fenced code
+  block, if one is present. If none are present it returns the full response.
+
+      llm 'JavaScript function for reversing a string' -x
+
 Options:
   -s, --system TEXT               System prompt to use
   -m, --model TEXT                Model to use
@@ -123,6 +128,7 @@ Options:
   --save TEXT                     Save prompt with this template name
   --async                         Run prompt asynchronously
   -u, --usage                     Show token usage
+  -x, --extract                   Extract first fenced code block
   --help                          Show this message and exit.
 ```
 

docs/usage.md

@@ -45,6 +45,24 @@ Some models support options. You can pass these using `-o/--option name value` -
 ```bash
 llm 'Ten names for cheesecakes' -o temperature 1.5
 ```
+
+(usage-extract-fenced-code)=
+### Extracting fenced code blocks
+
+If you are using an LLM to generate code it can be useful to retrieve just the code it produces without any of the surrounding explanatory text.
+
+The `-x/--extract` option will scan the response for the first instance of a Markdown fenced code block - something that looks like this:
+
+````
+```python
+def my_function():
+    # ...
+```
+````
+It will extract and returns just the content of that block, excluding the fenced coded delimiters. If there are no fenced code blocks it will return the full response.
+
+The entire response including explanatory text is still logged to the database, and can be viewed using `llm logs -c`.
+
 (usage-attachments)=
 ### Attachments
 

llm/cli.py

@@ -33,7 +33,12 @@ from llm import (
 
 from .migrations import migrate
 from .plugins import pm, load_plugins
-from .utils import mimetype_from_path, mimetype_from_string, token_usage_string
+from .utils import (
+    mimetype_from_path,
+    mimetype_from_string,
+    token_usage_string,
+    extract_first_fenced_code_block,
+)
 import base64
 import httpx
 import pathlib
@@ -204,6 +209,7 @@ def cli():
 @click.option("--save", help="Save prompt with this template name")
 @click.option("async_", "--async", is_flag=True, help="Run prompt asynchronously")
 @click.option("-u", "--usage", is_flag=True, help="Show token usage")
+@click.option("-x", "--extract", is_flag=True, help="Extract first fenced code block")
 def prompt(
     prompt,
     system,
@@ -222,6 +228,7 @@ def prompt(
     save,
     async_,
     usage,
+    extract,
 ):
     """
     Execute a prompt
@@ -243,12 +250,21 @@ def prompt(
         cat image | llm 'describe image' -a -
         # With an explicit mimetype:
         cat image | llm 'describe image' --at - image/jpeg
+
+    The -x/--extract option returns just the content of the first ``` fenced code
+    block, if one is present. If none are present it returns the full response.
+
+    \b
+        llm 'JavaScript function for reversing a string' -x
     """
     if log and no_log:
         raise click.ClickException("--log and --no-log are mutually exclusive")
 
     model_aliases = get_model_aliases()
 
+    if extract:
+        no_stream = True
+
     def read_prompt():
         nonlocal prompt
 
@@ -407,7 +423,10 @@ def prompt(
                         system=system,
                         **validated_options,
                     )
-                    print(await response.text())
+                    text = await response.text()
+                    if extract:
+                        text = extract_first_fenced_code_block(text) or text
+                    print(text)
                 return response
 
             response = asyncio.run(inner())
@@ -424,7 +443,10 @@ def prompt(
                     sys.stdout.flush()
                 print("")
             else:
-                print(response.text())
+                text = response.text()
+                if extract:
+                    text = extract_first_fenced_code_block(text) or text
+                print(text)
     except Exception as ex:
         raise click.ClickException(str(ex))
 

llm/default_plugins/openai_models.py

@@ -656,10 +656,11 @@ def combine_chunks(chunks: List) -> dict:
     }
     if logprobs:
         combined["logprobs"] = logprobs
-    for key in ("id", "object", "model", "created", "index"):
-        value = getattr(chunks[0], key, None)
-        if value is not None:
-            combined[key] = value
+    if chunks:
+        for key in ("id", "object", "model", "created", "index"):
+            value = getattr(chunks[0], key, None)
+            if value is not None:
+                combined[key] = value
 
     return combined
 

llm/utils.py

@@ -2,6 +2,7 @@ import click
 import httpx
 import json
 import puremagic
+import re
 import textwrap
 from typing import List, Dict, Optional
 
@@ -153,3 +154,38 @@ def token_usage_string(input_tokens, output_tokens, token_details) -> str:
     if token_details:
         bits.append(json.dumps(token_details))
     return ", ".join(bits)
+
+
+def extract_first_fenced_code_block(text: str) -> Optional[str]:
+    """
+    Extracts and returns the first Markdown fenced code block found in the given text.
+
+    The function handles fenced code blocks that:
+    - Use at least three backticks (`).
+    - May include a language tag immediately after the opening backticks.
+    - Use more than three backticks as long as the closing fence has the same number.
+
+    If no fenced code block is found, the function returns None.
+
+    Args:
+        text (str): The input text to search for a fenced code block.
+
+    Returns:
+        Optional[str]: The content of the first fenced code block, or None if not found.
+    """
+    # Regex pattern to match fenced code blocks
+    # - ^ or \n ensures that the fence is at the start of a line
+    # - (`{3,}) captures the opening backticks (at least three)
+    # - (\w+)? optionally captures the language tag
+    # - \n matches the newline after the opening fence
+    # - (.*?) non-greedy match for the code block content
+    # - \1 ensures that the closing fence has the same number of backticks
+    # - (?=\n|$) ensures that the closing fence is followed by a newline or end of string
+    pattern = re.compile(
+        r"""(?m)^(?P<fence>`{3,})(?P<lang>\w+)?\n(?P<code>.*?)^(?P=fence)(?=\n|$)""",
+        re.DOTALL,
+    )
+    match = pattern.search(text)
+    if match:
+        return match.group("code")
+    return None

tests/conftest.py

@@ -190,6 +190,27 @@ def mocked_openai_chat(httpx_mock):
     return httpx_mock
 
 
+@pytest.fixture
+def mocked_openai_chat_returning_fenced_code(httpx_mock):
+    httpx_mock.add_response(
+        method="POST",
+        url="https://api.openai.com/v1/chat/completions",
+        json={
+            "model": "gpt-4o-mini",
+            "usage": {},
+            "choices": [
+                {
+                    "message": {
+                        "content": "Code:\n\n````javascript\nfunction foo() {\n  return 'bar';\n}\n````\nDone.",
+                    }
+                }
+            ],
+        },
+        headers={"Content-Type": "application/json"},
+    )
+    return httpx_mock
+
+
 def stream_events():
     for delta, finish_reason in (
         ({"role": "assistant", "content": ""}, None),

tests/test_llm.py

@@ -322,6 +322,33 @@ def test_llm_default_prompt(
     )
 
 
+@pytest.mark.parametrize(
+    "args,expect_just_code",
+    (
+        (["-x"], True),
+        (["--extract"], True),
+        (["-x", "--async"], True),
+        (["--extract", "--async"], True),
+        # Use --no-stream here to ensure it passes test same as -x/--extract cases
+        (["--no-stream"], False),
+    ),
+)
+def test_extract_fenced_code(
+    mocked_openai_chat_returning_fenced_code, args, expect_just_code
+):
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        ["-m", "gpt-4o-mini", "--key", "x", "Write code"] + args,
+        catch_exceptions=False,
+    )
+    output = result.output
+    if expect_just_code:
+        assert "```" not in output
+    else:
+        assert "```" in output
+
+
 def test_openai_chat_stream(mocked_openai_chat_stream, user_path):
     runner = CliRunner()
     result = runner.invoke(cli, ["-m", "gpt-3.5-turbo", "--key", "x", "Say hi"])

tests/test_utils.py

@@ -1,5 +1,5 @@
 import pytest
-from llm.utils import simplify_usage_dict
+from llm.utils import simplify_usage_dict, extract_first_fenced_code_block
 
 
 @pytest.mark.parametrize(
@@ -40,3 +40,42 @@ from llm.utils import simplify_usage_dict
 )
 def test_simplify_usage_dict(input_data, expected_output):
     assert simplify_usage_dict(input_data) == expected_output
+
+
+@pytest.mark.parametrize(
+    "input,expected",
+    [
+        ["This is a sample text without any code blocks.", None],
+        [
+            "Here is some text.\n\n```\ndef foo():\n    return 'bar'\n```\n\nMore text.",
+            "def foo():\n    return 'bar'\n",
+        ],
+        [
+            "Here is some text.\n\n```python\ndef foo():\n    return 'bar'\n```\n\nMore text.",
+            "def foo():\n    return 'bar'\n",
+        ],
+        [
+            "Here is some text.\n\n````\ndef foo():\n    return 'bar'\n````\n\nMore text.",
+            "def foo():\n    return 'bar'\n",
+        ],
+        [
+            "Here is some text.\n\n````javascript\nfunction foo() {\n    return 'bar';\n}\n````\n\nMore text.",
+            "function foo() {\n    return 'bar';\n}\n",
+        ],
+        [
+            "Here is some text.\n\n```python\ndef foo():\n    return 'bar'\n````\n\nMore text.",
+            None,
+        ],
+        [
+            "First code block:\n\n```python\ndef foo():\n    return 'bar'\n```\n\nSecond code block:\n\n```javascript\nfunction foo() {\n    return 'bar';\n}\n```",
+            "def foo():\n    return 'bar'\n",
+        ],
+        [
+            "Here is some text.\n\n```python\ndef foo():\n    return `bar`\n```\n\nMore text.",
+            "def foo():\n    return `bar`\n",
+        ],
+    ],
+)
+def test_extract_first_fenced_code_block(input, expected):
+    actual = extract_first_fenced_code_block(input)
+    assert actual == expected