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: [PATCH] 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
---
 README.md         | 17 ++++++++++++++++-
 doc/tags          |  1 +
 doc/vim-ai.txt    | 19 +++++++++++++++++++
 py/utils.py       | 36 ++++++++++++++++++++++++++++++++++++
 syntax/aichat.vim |  1 +
 5 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index fa505dd..993e762 100644
--- a/README.md
+++ b/README.md
@@ -131,7 +131,22 @@ You are a Clean Code expert, I have the following code, please refactor it in a
 
 ```
 
-Supported chat roles are **`>>> system`**, **`>>> user`** and **`<<< assistant`**
+To include files in the chat a special `include` role is used:
+
+```
+>>> user
+
+Generate documentation for the following files
+
+>>> include
+
+/home/user/myproject/requirements.txt
+/home/user/myproject/**/*.py
+```
+
+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.gob` and relative paths to the current working directory (as determined by `getcwd()`) will be resolved to absolute paths.
+
+Supported chat roles are **`>>> system`**, **`>>> user`**, **`>>> include`** and **`<<< assistant`**
 
 ### `:AINewChat`
 
diff --git a/doc/tags b/doc/tags
index dfac716..45ad6dd 100644
--- a/doc/tags
+++ b/doc/tags
@@ -7,4 +7,5 @@ vim-ai	vim-ai.txt	/*vim-ai*
 vim-ai-about	vim-ai.txt	/*vim-ai-about*
 vim-ai-commands	vim-ai.txt	/*vim-ai-commands*
 vim-ai-config	vim-ai.txt	/*vim-ai-config*
+vim-ai-include	vim-ai.txt	/*vim-ai-include*
 vim-ai.txt	vim-ai.txt	/*vim-ai.txt*
diff --git a/doc/vim-ai.txt b/doc/vim-ai.txt
index f5f2cde..34765d7 100644
--- a/doc/vim-ai.txt
+++ b/doc/vim-ai.txt
@@ -104,6 +104,25 @@ Options: >
 Check OpenAI docs for more information:
 https://platform.openai.com/docs/api-reference/chat
 
+INCLUDE FILES                                  *vim-ai-include*
+
+To include files in the chat a special `include` role is used: >
+
+  >>> user
+
+  Generate documentation for the following files
+
+  >>> include
+
+  /home/user/myproject/requirements.txt
+  /home/user/myproject/**/*.py
+
+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.gob` and relative paths to the current
+working directory (as determined by `getcwd()`) will be resolved to absolute
+paths.
+
                                                 *:AINewChat*
 
 :AINewChat {preset shortname}?      spawn a new conversation with a given open
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():
diff --git a/syntax/aichat.vim b/syntax/aichat.vim
index ecd8459..1f32efb 100644
--- a/syntax/aichat.vim
+++ b/syntax/aichat.vim
@@ -1,5 +1,6 @@
 syntax match aichatRole ">>> system"
 syntax match aichatRole ">>> user"
+syntax match aichatRole ">>> include"
 syntax match aichatRole "<<< assistant"
 
 highlight default link aichatRole Comment