diff --git a/README.md b/README.md index 1dcee9860..8ca76c715 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt. Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > Tell me a joke about a pelican Why don't pelicans like to tip waiters? diff --git a/docs/fragments.md b/docs/fragments.md index 366505b55..235f9736c 100644 --- a/docs/fragments.md +++ b/docs/fragments.md @@ -58,6 +58,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt. Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > Explain this document to me ``` @@ -69,6 +70,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt. Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > !fragment https://llm.datasette.io/en/stable/fragments.html ``` diff --git a/docs/index.md b/docs/index.md index 939ef3807..2bee2a460 100644 --- a/docs/index.md +++ b/docs/index.md @@ -96,6 +96,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt. Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > Tell me a joke about a pelican Why don't pelicans like to tip waiters? diff --git a/docs/plugins/plugin-hooks.md b/docs/plugins/plugin-hooks.md index 4b8b2922a..e7ba36678 100644 --- a/docs/plugins/plugin-hooks.md +++ b/docs/plugins/plugin-hooks.md @@ -141,6 +141,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > Remember my name is Henry Tool call: Memory_set({'key': 'user_name', 'value': 'Henry'}) diff --git a/docs/usage.md b/docs/usage.md index f4271cddf..40814aefb 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -518,6 +518,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > who are you? I am a sentient cheesecake, meaning I am an artificial intelligence embodied in a dessert form, specifically a @@ -540,6 +541,7 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt. Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > !multi custom-end Explain this error: @@ -560,9 +562,27 @@ Type 'exit' or 'quit' to exit Type '!multi' to enter multiple lines, then '!end' to finish Type '!edit' to open your default editor and modify the prompt. Type '!fragment [ ...]' to insert one or more fragments +Type '!attach ' to attach a file > !edit ``` +For multi-modal models that support attachments, use `!attach ` to queue a file or URL for the next message: + +```bash +> !attach photo.jpg +Attachment queued for next message +> Describe this image +``` + +You can queue multiple attachments before entering your prompt. `!attach` can also be combined with `!multi`: + +```bash +> !multi +Describe this image +!attach photo.jpg +!end +``` + `llm chat` takes the same `--tool/-T` and `--functions` options as `llm prompt`. You can use this to start a chat with the specified {ref}`tools ` enabled. ## Listing available models diff --git a/llm/cli.py b/llm/cli.py index a9968160d..23f7d8383 100644 --- a/llm/cli.py +++ b/llm/cli.py @@ -153,6 +153,7 @@ def _run_chat( click.echo( "Type '!fragment [ ...]' to insert one or more fragments" ) + click.echo("Type '!attach ' to attach a file") argument_fragments = list(initial_fragments or []) argument_attachments = list(initial_attachments or []) @@ -166,13 +167,6 @@ def _run_chat( prompt = click.prompt("", prompt_suffix="> " if not in_multi else "") fragments = [] attachments = [] - if argument_fragments: - fragments += argument_fragments - # Fragments from command options are added to the first message only. - argument_fragments = [] - if argument_attachments: - attachments = argument_attachments - argument_attachments = [] if prompt.strip().startswith("!multi"): in_multi = True bits = prompt.strip().split() @@ -186,7 +180,26 @@ def _run_chat( continue prompt = edited_prompt.strip() if db is not None and prompt.strip().startswith("!fragment "): - prompt, fragments, attachments = process_fragments_in_chat(db, prompt) + prompt, extra_fragments, fragment_attachments = process_fragments_in_chat( + db, prompt + ) + fragments += extra_fragments + attachments += fragment_attachments + + try: + prompt, extra_attachments = process_attachments_in_chat(prompt) + except click.ClickException as ex: + # Invalid attachments should not terminate an interactive session. + click.echo(f"Error: {ex.format_message()}", err=True) + continue + if extra_attachments and not in_multi and not prompt.strip(): + # A standalone !attach command queues attachments for the next + # actual message instead of sending an empty prompt immediately. + argument_attachments.extend(extra_attachments) + noun = "Attachment" if len(extra_attachments) == 1 else "Attachments" + click.echo(f"{noun} queued for next message") + continue + attachments += extra_attachments if in_multi: if prompt.strip() == end_token: @@ -209,6 +222,13 @@ def _run_chat( if transform_prompt is not None: prompt = transform_prompt(prompt) + # Command-line fragments and attachments, plus standalone !attach + # commands, remain pending until an actual message is sent. + fragments = argument_fragments + fragments + attachments = argument_attachments + attachments + argument_fragments = [] + argument_attachments = [] + response = prompt_callback(prompt, fragments, attachments) display_stream_events( response.stream_events(), @@ -329,6 +349,29 @@ def process_fragments_in_chat( return "\n".join(prompt_lines), fragments, attachments +def process_attachments_in_chat(prompt: str) -> tuple[str, list[Attachment]]: + """ + Process any !attach commands in a chat prompt and return the modified prompt plus resolved attachments. + """ + prompt_lines = [] + attachments = [] + for line in prompt.splitlines(): + stripped_line = line.strip() + if stripped_line == "!attach": + raise click.ClickException("Usage: !attach ") + if stripped_line.startswith("!attach "): + attachment_value = stripped_line.removeprefix("!attach ").strip() + if not attachment_value: + raise click.ClickException("Usage: !attach ") + try: + attachments.append(resolve_attachment(attachment_value)) + except AttachmentError as ex: + raise click.ClickException(str(ex)) + else: + prompt_lines.append(line) + return "\n".join(prompt_lines), attachments + + class AttachmentError(Exception): """Exception raised for errors in attachment resolution.""" diff --git a/tests/test_chat.py b/tests/test_chat.py index 120f757f7..59f04c1ba 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -32,6 +32,17 @@ def logged_rows(db): ] +TINY_PNG = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\xa6\x00\x00\x01\x1a" + b"\x02\x03\x00\x00\x00\xe6\x99\xc4^\x00\x00\x00\tPLTE\xff\xff\xff" + b"\x00\xff\x00\xfe\x01\x00\x12t\x01J\x00\x00\x00GIDATx\xda\xed\xd81\x11" + b"\x000\x08\xc0\xc0.]\xea\xaf&Q\x89\x04V\xe0>\xf3+\xc8\x91Z\xf4\xa2\x08EQ\x14E" + b"Q\x14EQ\x14EQ\xd4B\x91$I3\xbb\xbf\x08EQ\x14EQ\x14EQ\x14E\xd1\xa5" + b"\xd4\x17\x91\xc6\x95\x05\x15\x0f\x9f\xc5\t\x9f\xa4\x00\x00\x00\x00IEND\xaeB`" + b"\x82" +) + + @pytest.mark.xfail(sys.platform == "win32", reason="Expected to fail on Windows") def test_chat_basic(mock_model, logs_db): runner = CliRunner() @@ -50,6 +61,7 @@ def test_chat_basic(mock_model, logs_db): "\nType '!multi' to enter multiple lines, then '!end' to finish" "\nType '!edit' to open your default editor and modify the prompt" "\nType '!fragment [ ...]' to insert one or more fragments" + "\nType '!attach ' to attach a file" "\n> Hi" "\none world" "\n> Hi two" @@ -99,6 +111,7 @@ def test_chat_basic(mock_model, logs_db): "\nType '!multi' to enter multiple lines, then '!end' to finish" "\nType '!edit' to open your default editor and modify the prompt" "\nType '!fragment [ ...]' to insert one or more fragments" + "\nType '!attach ' to attach a file" "\n> Continue" "\ncontinued" "\n> quit" @@ -135,6 +148,7 @@ def test_chat_system(mock_model, logs_db): "\nType '!multi' to enter multiple lines, then '!end' to finish" "\nType '!edit' to open your default editor and modify the prompt" "\nType '!fragment [ ...]' to insert one or more fragments" + "\nType '!attach ' to attach a file" "\n> Hi" "\nI am mean" "\n> quit" @@ -307,6 +321,7 @@ def upper(text: str) -> str: "Type '!multi' to enter multiple lines, then '!end' to finish\n" "Type '!edit' to open your default editor and modify the prompt\n" "Type '!fragment [ ...]' to insert one or more fragments\n" + "Type '!attach ' to attach a file\n" '> {"prompt": "Convert hello to uppercase", "tool_calls": [{"name": "upper", ' '"arguments": {"text": "hello"}}]}\n' "{\n" @@ -355,3 +370,102 @@ def test_chat_fragments(tmpdir): ).output assert '"prompt": "one' in output assert '"prompt": "two"' in output + + +@pytest.mark.xfail(sys.platform == "win32", reason="Expected to fail on Windows") +def test_chat_attach_queued_for_next_prompt(tmp_path, mock_model, logs_db): + image_path = tmp_path / "image with spaces.png" + image_path.write_bytes(TINY_PNG) + second_image_path = tmp_path / "second image.png" + second_image_path.write_bytes(TINY_PNG) + missing_path = tmp_path / "missing image.png" + runner = CliRunner() + mock_model.enqueue(["saw image"]) + mock_model.enqueue(["follow-up"]) + result = runner.invoke( + llm.cli.cli, + ["chat", "-m", "mock"], + input=( + f"!attach {image_path}\n" + f"!attach {second_image_path}\n" + f"!attach {missing_path}\n" + "Describe this image\n" + "What did I just ask you to do?\n" + "quit\n" + ), + catch_exceptions=False, + ) + assert result.exit_code == 0 + assert result.output.endswith("\n> quit\n") + assert "Attachment queued for next message" in result.output + assert f"Error: File {missing_path} does not exist" in result.output + assert result.output.index("Error: File") < result.output.index( + "> Describe this image" + ) + assert len(mock_model.history) == 2 + prompt = mock_model.history[0][0] + assert prompt.prompt == "Describe this image" + assert prompt.attachments == [ + llm.Attachment( + type="image/png", + path=str(image_path), + url=None, + content=None, + _id=ANY, + ), + llm.Attachment( + type="image/png", + path=str(second_image_path), + url=None, + content=None, + _id=ANY, + ), + ] + follow_up_prompt = mock_model.history[1][0] + assert follow_up_prompt.prompt == "What did I just ask you to do?" + assert follow_up_prompt.attachments == [] + + +@pytest.mark.xfail(sys.platform == "win32", reason="Expected to fail on Windows") +def test_chat_multi_fragments_and_attach(tmp_path, mock_model, logs_db): + fragment_path = tmp_path / "fragment.txt" + fragment_path.write_text("fragment text") + image_path = tmp_path / "image.png" + image_path.write_bytes(TINY_PNG) + runner = CliRunner() + mock_model.enqueue(["parsed"]) + result = runner.invoke( + llm.cli.cli, + ["chat", "-m", "mock"], + input=( + "!multi\n" + "Describe this image using this fragment:\n" + f"!fragment {fragment_path}\n" + f"!attach {image_path}\n" + "!end\n" + "quit\n" + ), + catch_exceptions=False, + ) + assert result.exit_code == 0 + assert result.output.endswith("\n> quit\n") + prompt = mock_model.history[0][0] + assert prompt.prompt == "fragment text\nDescribe this image using this fragment:" + assert prompt.fragments == ["fragment text"] + assert prompt.attachments == [ + llm.Attachment( + type="image/png", + path=str(image_path), + url=None, + content=None, + _id=ANY, + ) + ] + attachment = next(iter(logs_db["attachments"].rows)) + assert attachment == { + "id": ANY, + "type": "image/png", + "path": str(image_path), + "url": None, + "content": None, + }