diff options
| author | Jason Kölker <jason@koelker.net> | 2024-01-24 02:00:07 +0000 |
|---|---|---|
| committer | Jason Kölker <jason@koelker.net> | 2024-01-24 02:34:44 +0000 |
| commit | 9e6f8e8d037a477c55939ece36ab7b000f876e85 (patch) | |
| tree | 457b7c81b511b42ece0704e03b673702a6786d0e /py | |
| parent | 4692eec84b5aa9d95256bef515bd1d17471e5570 (diff) | |
| download | vim-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.py | 36 |
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(): |