feat: support paths/URLs in any prompt, refactored entrypoint to call…

… a new public API with core logic (#37)

* feat: refactored CLI into calling a public API surface, also added support for reading paths anywhere in prompt

* fix: fix reading files in prompt, added test for passing URLs on CLI

* test: skip test_url if playwright not installed

* fix: fixed handling paths in prompts

gptme/cli.py

@@ -4,6 +4,7 @@
 import io
 import logging
 import os
+import re
 import readline  # noqa: F401
 import sys
 import urllib.parse
@@ -52,9 +53,15 @@
 {action_readme}"""
 
 
-def init(verbose: bool, llm: LLMChoice, model: str, interactive: bool):
-    # log init
-    logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO)
+_init_done = False
+
+
+def init(llm: LLMChoice, model: str, interactive: bool):
+    global _init_done
+    if _init_done:
+        logger.warning("init() called twice, ignoring")
+        return
+    _init_done = True
 
     # init
     logger.debug("Started")
@@ -73,6 +80,11 @@ def init(verbose: bool, llm: LLMChoice, model: str, interactive: bool):
     init_tools()
 
 
+def init_logging(verbose):
+    # log init
+    logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO)
+
+
 @click.command(help=docstring)
 @click.argument("prompts", default=None, required=False, nargs=-1)
 @click.option(
@@ -144,7 +156,8 @@ def main(
     if "PYTEST_CURRENT_TEST" in os.environ:
         interactive = False
 
-    init(verbose, llm, model, interactive)
+    # init logging
+    init_logging(verbose)
 
     if not interactive:
         no_confirm = True
@@ -153,14 +166,14 @@ def main(
         logger.warning("Skipping all confirmation prompts.")
 
     # get initial system prompt
-    prompt_msgs = [get_prompt(prompt_system)]
+    initial_msgs = [get_prompt(prompt_system)]
 
     # if stdin is not a tty, we're getting piped input, which we should include in the prompt
     if not sys.stdin.isatty():
         # fetch prompt from stdin
         prompt_stdin = _read_stdin()
         if prompt_stdin:
-            prompt_msgs += [Message("system", f"```stdin\n{prompt_stdin}\n```")]
+            initial_msgs += [Message("system", f"```stdin\n{prompt_stdin}\n```")]
 
             # Attempt to switch to interactive mode
             sys.stdin.close()
@@ -172,31 +185,60 @@ def main(
                     "Failed to switch to interactive mode, continuing in non-interactive mode"
                 )
 
+    # join prompts, grouped by `-` if present, since that's the separator for multiple-round prompts
+    sep = "\n\n" + MULTIPROMPT_SEPARATOR
+    prompts = [p.strip() for p in "\n\n".join(prompts).split(sep) if p]
+    prompt_msgs = [Message("user", p) for p in prompts]
+
+    chat(
+        prompt_msgs,
+        initial_msgs,
+        name,
+        llm,
+        model,
+        stream,
+        no_confirm,
+        interactive,
+        show_hidden,
+    )
+
+
+def chat(
+    prompt_msgs: list[Message],
+    initial_msgs: list[Message],
+    name: str,
+    llm: LLMChoice,
+    model: ModelChoice,
+    stream: bool = True,
+    no_confirm: bool = False,
+    interactive: bool = True,
+    show_hidden: bool = False,
+):
+    """
+    Run the chat loop.
+
+    Callable from other modules.
+    """
+    # init
+    init(llm, model, interactive)
+
     # we need to run this before checking stdin, since the interactive doesn't work with the switch back to interactive mode
     logfile = get_logfile(
-        name, interactive=(not prompts and interactive) and sys.stdin.isatty()
+        name, interactive=(not prompt_msgs and interactive) and sys.stdin.isatty()
     )
     print(f"Using logdir {logfile.parent}")
-    log = LogManager.load(logfile, initial_msgs=prompt_msgs, show_hidden=show_hidden)
+    log = LogManager.load(logfile, initial_msgs=initial_msgs, show_hidden=show_hidden)
 
     # print log
     log.print()
     print("--- ^^^ past messages ^^^ ---")
 
-    # check if any prompt is a full path, if so, replace it with the contents of that file
-    # TODO: add support for directories
-    # TODO: maybe do this for all prompts, not just those passed on cli
-    prompts = [_parse_prompt(p) for p in prompts]
-    # join prompts, grouped by `-` if present, since that's the separator for multiple-round prompts
-    sep = "\n\n" + MULTIPROMPT_SEPARATOR
-    prompts = [p.strip() for p in "\n\n".join(prompts).split(sep) if p]
-
     # main loop
     while True:
-        # if prompts given on cli, insert next prompt into log
-        if prompts:
-            prompt = prompts.pop(0)
-            msg = Message("user", prompt)
+        # if prompt_msgs given, insert next prompt into log
+        if prompt_msgs:
+            msg = prompt_msgs.pop(0)
+            msg = _include_paths(msg)
             log.append(msg)
             # if prompt is a user-command, execute it
             if execute_cmd(msg, log):
@@ -251,7 +293,9 @@ def loop(
             # Empty command, ask for input again
             print()
             return
-        yield Message("user", inquiry, quiet=True)
+        msg = Message("user", inquiry, quiet=True)
+        msg = _include_paths(msg)
+        yield msg
 
     # print response
     try:
@@ -310,6 +354,7 @@ def get_name(name: str) -> Path:
 
 
 # default history if none found
+# NOTE: there are also good examples in the integration tests
 history_examples = [
     "What is love?",
     "Have you heard about an open-source app called ActivityWatch?",
@@ -420,13 +465,60 @@ def _read_stdin() -> str:
     return all_data
 
 
-def _parse_prompt(prompt: str) -> str:
+def _include_paths(msg: Message) -> Message:
+    """Searches the message for any valid paths and appends the contents of such files as codeblocks."""
+    # TODO: add support for directories?
+    assert msg.role == "user"
+
+    # list the current directory
+    cwd_files = [f.name for f in Path.cwd().iterdir()]
+
+    # match absolute, home, relative paths, and URLs anywhere in the message
+    # could be wrapped with spaces or backticks, possibly followed by a question mark
+    # don't look in codeblocks, and don't match paths that are already in codeblocks
+    # TODO: this will misbehave if there are codeblocks (or triple backticks) in codeblocks
+    content_no_codeblocks = re.sub(r"```.*?\n```", "", msg.content, flags=re.DOTALL)
+    append_msg = ""
+    for word in re.split(r"[\s`]", content_no_codeblocks):
+        # remove wrapping backticks
+        word = word.strip("`")
+        # remove trailing question mark
+        word = word.rstrip("?")
+        if not word:
+            continue
+        if (
+            # if word starts with a path character
+            word[0] in ["/", "~", "."]
+            # or word is a URL
+            or word.startswith("http")
+            # or word is a file in the current dir,
+            # or a path that starts in a folder in the current dir
+            or any(word.split("/", 1)[0] == file for file in cwd_files)
+        ):
+            logger.debug(f"potential path/url: {word=}")
+            p = _parse_prompt(word)
+            if p:
+                # if we found a valid path, replace it with the contents of the file
+                append_msg += "\n\n" + p
+
+    # append the message with the file contents
+    if append_msg:
+        msg.content += append_msg
+
+    return msg
+
+
+def _parse_prompt(prompt: str) -> str | None:
+    """
+    Takes a string that might be a path,
+    and if so, returns the contents of that file wrapped in a codeblock.
+    """
     # if prompt is a command, exit early (as commands might take paths as arguments)
     if any(
         prompt.startswith(command)
         for command in [f"{CMDFIX}{cmd}" for cmd in action_descriptions.keys()]
     ):
-        return prompt
+        return None
 
     try:
         # check if prompt is a path, if so, replace it with the contents of that file
@@ -458,10 +550,17 @@ def _parse_prompt(prompt: str) -> str:
                 urls.append(word)
         except ValueError:
             pass
+
+    result = ""
     if paths or urls:
-        prompt += "\n\n"
+        result += "\n\n"
+        if paths:
+            logger.debug(f"{paths=}")
+        if urls:
+            logger.debug(f"{urls=}")
     for path in paths:
-        prompt += _parse_prompt(path)
+        result += _parse_prompt(path) or ""
+
     for url in urls:
         try:
             # noreorder
@@ -475,11 +574,11 @@ def _parse_prompt(prompt: str) -> str:
 
         try:
             content = read_url(url)
-            prompt += f"```{url}\n{content}\n```"
+            result += f"```{url}\n{content}\n```"
         except Exception as e:
             logger.warning(f"Failed to read URL {url}: {e}")
 
-    return prompt
+    return result
 
 
 if __name__ == "__main__":

gptme/cli_server.py

@@ -2,7 +2,7 @@
 
 from gptme.util import logger
 
-from .cli import init
+from .cli import init, init_logging
 
 
 @click.command("gptme-server")
@@ -20,7 +20,8 @@
 )
 def main(verbose, llm, model):  # pragma: no cover
     """Starts a server and web UI."""
-    init(verbose, llm, model, interactive=False)
+    init_logging(verbose)
+    init(llm, model, interactive=False)
 
     # if flask not installed, ask the user to install `server` extras
     try:

gptme/prompts.py

@@ -147,7 +147,7 @@ def prompt_tools() -> Generator[Message, None, None]:
 
 ## python
 
-When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment. 
+When you send a message containing Python code (and is not a file block), it will be executed in a stateful environment.
 Python will respond with the output of the execution.
 
 The following libraries are available:

tests/test_cli.py

@@ -1,3 +1,4 @@
+import importlib
 import os
 import random
 from pathlib import Path
@@ -249,6 +250,18 @@ def test_stdin(args: list[str], runner: CliRunner):
     assert result.exit_code == 0
 
 
+@pytest.mark.slow
+@pytest.mark.skipif(
+    importlib.util.find_spec("playwright") is None,
+    reason="playwright not installed",
+)
+def test_url(args: list[str], runner: CliRunner):
+    args.append("Who is the CEO of https://superuserlabs.org?")
+    result = runner.invoke(gptme.cli.main, args)
+    assert "Erik Bjäreholt" in result.output
+    assert result.exit_code == 0
+
+
 def test_version(args: list[str], runner: CliRunner):
     args.append("--version")
     result = runner.invoke(gptme.cli.main, args)

tests/test_server.py

@@ -14,7 +14,7 @@
 
 @pytest.fixture(autouse=True)
 def init_():
-    init(verbose=False, llm="openai", model="gpt-3.5-turbo", interactive=False)
+    init(llm="openai", model="gpt-3.5-turbo", interactive=False)
 
 
 @pytest.fixture