From 9e6f8e8d037a477c55939ece36ab7b000f876e85 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jason=20K=C3=B6lker?= <jason@koelker.net>
Date: Wed, 24 Jan 2024 02:00:07 +0000
Subject: 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
---
 py/utils.py | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

(limited to 'py/utils.py')

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():
-- 
cgit v1.2.3