From 639e4cb1739c22997e0bf34e8bf0856741e45412 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Wed, 6 Nov 2024 17:22:06 +0100
Subject: [PATCH] refactor: move get_workspace_prompt into prompts.py, add TODO
 comments

---
 gptme/chat.py    |  6 ++++--
 gptme/config.py  | 26 +++++---------------------
 gptme/prompts.py | 25 ++++++++++++++++++++++++-
 3 files changed, 33 insertions(+), 24 deletions(-)

diff --git a/gptme/chat.py b/gptme/chat.py
index 35b1be8f..598b7777 100644
--- a/gptme/chat.py
+++ b/gptme/chat.py
@@ -9,7 +9,6 @@
 from pathlib import Path
 
 from .commands import action_descriptions, execute_cmd
-from .config import get_workspace_prompt
 from .constants import PROMPT_USER
 from .init import init
 from .interrupt import clear_interruptible, set_interruptible
@@ -17,6 +16,7 @@
 from .logmanager import Log, LogManager, prepare_messages
 from .message import Message
 from .models import get_model
+from .prompts import get_workspace_prompt
 from .readline import add_history
 from .tools import ToolUse, execute_msg, has_tool
 from .tools.base import ConfirmFunc
@@ -78,7 +78,9 @@ def chat(
     console.log(f"Using workspace at {path_with_tilde(workspace)}")
     os.chdir(workspace)
 
-    workspace_prompt = get_workspace_prompt(str(workspace))
+    workspace_prompt = get_workspace_prompt(workspace)
+    # FIXME: this is hacky
+    # NOTE: needs to run after the workspace is set
     # check if message is already in log, such as upon resume
     if (
         workspace_prompt
diff --git a/gptme/config.py b/gptme/config.py
index 7aedfc97..4002a43d 100644
--- a/gptme/config.py
+++ b/gptme/config.py
@@ -1,4 +1,3 @@
-import glob
 import logging
 import os
 from dataclasses import dataclass, field
@@ -124,12 +123,12 @@ def set_config_value(key: str, value: str) -> None:  # pragma: no cover
     _config = load_config()
 
 
-def get_workspace_prompt(workspace: str) -> str:
+def get_project_config(workspace: Path) -> ProjectConfig | None:
     project_config_paths = [
         p
         for p in (
-            Path(workspace) / "gptme.toml",
-            Path(workspace) / ".github" / "gptme.toml",
+            workspace / "gptme.toml",
+            workspace / ".github" / "gptme.toml",
         )
         if p.exists()
     ]
@@ -141,23 +140,8 @@ def get_workspace_prompt(workspace: str) -> str:
         # load project config
         with open(project_config_path) as f:
             project_config = tomlkit.load(f)
-            project = ProjectConfig(**project_config)  # type: ignore
-            files = []
-            for file in project.files:
-                # expand user
-                file = str(Path(file).expanduser())
-                # expand with glob
-                if new_files := glob.glob(file):
-                    files.extend(new_files)
-                else:
-                    logger.error(
-                        f"File {file} specified in project config does not exist"
-                    )
-                    exit(1)
-        return "\n\nSelected project files, read more with cat:\n" + "\n\n".join(
-            [f"```{Path(file).name}\n{Path(file).read_text()}\n```" for file in files]
-        )
-    return ""
+        return ProjectConfig(**project_config)  # type: ignore
+    return None
 
 
 if __name__ == "__main__":
diff --git a/gptme/prompts.py b/gptme/prompts.py
index a811dbcf..027f5341 100644
--- a/gptme/prompts.py
+++ b/gptme/prompts.py
@@ -5,15 +5,17 @@
 When prompting, it is important to provide clear instructions and avoid any ambiguity.
 """
 
+import glob
 import logging
 import os
 import platform
 import subprocess
 from collections.abc import Generator, Iterable
+from pathlib import Path
 from typing import Literal
 
 from .__version__ import __version__
-from .config import get_config
+from .config import get_config, get_project_config
 from .message import Message
 from .tools import loaded_tools
 from .util import document_prompt_function
@@ -235,6 +237,27 @@ def prompt_systeminfo() -> Generator[Message, None, None]:
     )
 
 
+def get_workspace_prompt(workspace: Path) -> str:
+    # NOTE: needs to run after the workspace is initialized (i.e. initial prompt is constructed)
+    # TODO: update this prompt if the files change
+    # TODO: include `git status/diff/log` summary, and keep it up-to-date
+    if project := get_project_config(workspace):
+        files = []
+        for file in project.files:
+            # expand user
+            file = str(Path(file).expanduser())
+            # expand with glob
+            if new_files := glob.glob(file):
+                files.extend(new_files)
+            else:
+                logger.error(f"File {file} specified in project config does not exist")
+                exit(1)
+        return "\n\nSelected project files, read more with cat:\n" + "\n\n".join(
+            [f"```{Path(file).name}\n{Path(file).read_text()}\n```" for file in files]
+        )
+    return ""
+
+
 document_prompt_function(interactive=True)(prompt_gptme)
 document_prompt_function()(prompt_user)
 document_prompt_function()(prompt_project)