summaryrefslogtreecommitdiff
path: root/py/utils.py
diff options
context:
space:
mode:
authorJason Kölker <jason@koelker.net>2024-01-24 02:00:07 +0000
committerJason Kölker <jason@koelker.net>2024-01-24 02:34:44 +0000
commit9e6f8e8d037a477c55939ece36ab7b000f876e85 (patch)
tree457b7c81b511b42ece0704e03b673702a6786d0e /py/utils.py
parent4692eec84b5aa9d95256bef515bd1d17471e5570 (diff)
downloadvim-ai-9e6f8e8d037a477c55939ece36ab7b000f876e85.tar.gz
feat(chat): add `include` role to include files
Files may be included in the chat by a special `include` role. Each file's contents will be added to an additional `user` role message with the files separated by `==> {path} <==` where `{path}` is the path to the file. Globbing is expanded out via `glob.glob` and relative apths to the current working directory (as determined by `getcwd()`) will be resolved to absolute paths. Example: ``` >>> user Generate documentation for the following files >>> include /home/user/myproject/src/../requirements.txt /home/user/myproject/**/*.py ``` Fixes: #69
Diffstat (limited to '')
-rw-r--r--py/utils.py36
1 files changed, 36 insertions, 0 deletions
diff --git a/py/utils.py b/py/utils.py
index d04b18d..a5c489f 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -1,5 +1,6 @@
import vim
import datetime
+import glob
import sys
import os
import json
@@ -103,6 +104,7 @@ def render_text_chunks(chunks, is_selection):
if not full_text.strip():
print_info_message('Empty response received. Tip: You can try modifying the prompt and retry.')
+
def parse_chat_messages(chat_content):
lines = chat_content.splitlines()
messages = []
@@ -113,6 +115,9 @@ def parse_chat_messages(chat_content):
if line.startswith(">>> user"):
messages.append({"role": "user", "content": ""})
continue
+ if line.startswith(">>> include"):
+ messages.append({"role": "include", "content": ""})
+ continue
if line.startswith("<<< assistant"):
messages.append({"role": "assistant", "content": ""})
continue
@@ -124,6 +129,37 @@ def parse_chat_messages(chat_content):
# strip newlines from the content as it causes empty responses
message["content"] = message["content"].strip()
+ if message["role"] == "include":
+ message["role"] = "user"
+ paths = message["content"].split("\n")
+ message["content"] = ""
+
+ pwd = vim.eval("getcwd()")
+ for i in range(len(paths)):
+ path = os.path.expanduser(paths[i])
+ if not os.path.isabs(path):
+ path = os.path.join(pwd, path)
+
+ paths[i] = path
+
+ if '**' in path:
+ paths[i] = None
+ paths.extend(glob.glob(path, recursive=True))
+
+ for path in paths:
+ if path is None:
+ continue
+
+ if os.path.isdir(path):
+ continue
+
+ try:
+ with open(path, "r") as file:
+ message["content"] += f"\n\n==> {path} <==\n" + file.read()
+ except UnicodeDecodeError:
+ message["content"] += "\n\n" + f"==> {path} <=="
+ message["content"] += "\n" + "Binary file, cannot display"
+
return messages
def parse_chat_header_options():