|
1
|
|
|
import sys |
|
2
|
|
|
|
|
3
|
|
|
import rich_click as click |
|
4
|
|
|
from click import Abort |
|
5
|
|
|
from rich.console import Console |
|
6
|
|
|
from rich.progress import Progress |
|
7
|
|
|
from rich.prompt import Confirm |
|
8
|
|
|
|
|
9
|
|
|
from menderbot import __version__ |
|
10
|
|
|
from menderbot.check import run_check |
|
11
|
|
|
from menderbot.config import create_default_config, has_config, has_llm_consent |
|
12
|
|
|
from menderbot.git_client import git_commit, git_diff_head, git_show_top_level |
|
13
|
|
|
from menderbot.ingest import ask_index, get_chat_engine, index_exists, ingest_repo |
|
14
|
|
|
from menderbot.llm import ( |
|
15
|
|
|
INSTRUCTIONS, |
|
16
|
|
|
get_response, |
|
17
|
|
|
has_key, |
|
18
|
|
|
key_env_var, |
|
19
|
|
|
unwrap_codeblock, |
|
20
|
|
|
) |
|
21
|
|
|
from menderbot.prompts import ( |
|
22
|
|
|
change_list_prompt, |
|
23
|
|
|
code_review_prompt, |
|
24
|
|
|
commit_msg_prompt, |
|
25
|
|
|
type_prompt, |
|
26
|
|
|
) |
|
27
|
|
|
from menderbot.source_file import SourceFile |
|
28
|
|
|
|
|
29
|
|
|
console = Console() |
|
30
|
|
|
|
|
31
|
|
|
|
|
32
|
|
|
@click.group(context_settings=dict(help_option_names=["-h", "--help"])) |
|
33
|
|
|
@click.version_option(__version__, prog_name="menderbot") |
|
34
|
|
|
@click.pass_context |
|
35
|
|
|
def cli(ctx): |
|
36
|
|
|
""" |
|
37
|
|
|
An AI-powered command line tool for working with legacy code. |
|
38
|
|
|
|
|
39
|
|
|
You can try using --help at the top level and also for |
|
40
|
|
|
specific subcommands. |
|
41
|
|
|
|
|
42
|
|
|
Connects to OpenAI using OPENAI_API_KEY environment variable. |
|
43
|
|
|
""" |
|
44
|
|
|
if not has_key(): |
|
45
|
|
|
console.log(f"{key_env_var()} not found in env, will not be able to connect.") |
|
46
|
|
|
ctx.ensure_object(dict) |
|
47
|
|
|
|
|
48
|
|
|
|
|
49
|
|
|
@cli.command() |
|
50
|
|
|
@click.argument("q", required=False) |
|
51
|
|
|
def ask(q): |
|
52
|
|
|
"""Ask a question about a specific piece of code or concept.""" |
|
53
|
|
|
if not index_exists(): |
|
54
|
|
|
console.print("[red]Index not found[/red]: please run menderbot ingest") |
|
55
|
|
|
return |
|
56
|
|
|
check_llm_consent() |
|
57
|
|
|
new_question = q |
|
58
|
|
|
if not new_question: |
|
59
|
|
|
new_question = console.input("[green]Ask[/green]: ") |
|
60
|
|
|
with Progress(transient=True) as progress: |
|
61
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
62
|
|
|
response = ask_index(new_question) |
|
63
|
|
|
progress.update(task, completed=True) |
|
64
|
|
|
console.print(f"[cyan]Bot[/cyan]: {response}") |
|
65
|
|
|
|
|
66
|
|
|
|
|
67
|
|
|
@cli.command() |
|
68
|
|
|
def chat(): |
|
69
|
|
|
"""Interactively chat in the context of the current directory.""" |
|
70
|
|
|
if not index_exists(): |
|
71
|
|
|
console.print("[red]Index not found[/red]: please run menderbot ingest") |
|
72
|
|
|
else: |
|
73
|
|
|
console.print("Loading index...") |
|
74
|
|
|
check_llm_consent() |
|
75
|
|
|
chat_engine = get_chat_engine() |
|
76
|
|
|
while True: |
|
77
|
|
|
new_question = console.input("[green]Ask[/green]: ") |
|
78
|
|
|
# new_question += "\nUse your tool to query for context." |
|
79
|
|
|
if new_question: |
|
80
|
|
|
streaming_response = chat_engine.stream_chat(new_question) |
|
81
|
|
|
console.print("[cyan]Bot[/cyan]: ", end="") |
|
82
|
|
|
for token in streaming_response.response_gen: |
|
83
|
|
|
console.out(token, end="") |
|
84
|
|
|
console.out("\n") |
|
85
|
|
|
|
|
86
|
|
|
|
|
87
|
|
|
def try_function_type_hints(mypy_cmd, source_file, function_ast, needs_typing): |
|
88
|
|
|
from menderbot.typing import add_type_hints, parse_type_hint_answer # Lazy import |
|
89
|
|
|
|
|
90
|
|
|
function_text = function_ast.text |
|
91
|
|
|
check_command = ( |
|
92
|
|
|
f"{mypy_cmd} --shadow-file {source_file.path} {source_file.path}.shadow" |
|
93
|
|
|
) |
|
94
|
|
|
max_tries = 2 |
|
95
|
|
|
check_output = None |
|
96
|
|
|
# First set them all to wrong type, to produce an error message. |
|
97
|
|
|
hints = [(ident, "None") for ident in needs_typing] |
|
98
|
|
|
# TODO |
|
99
|
|
|
imports = [] |
|
100
|
|
|
insertions_for_function = add_type_hints(function_ast, hints, imports) |
|
101
|
|
|
if insertions_for_function: |
|
102
|
|
|
source_file.update_file(insertions_for_function, suffix=".shadow") |
|
103
|
|
|
console.print("> ", check_command) |
|
104
|
|
|
(success, pre_check_output) = run_check(check_command) |
|
105
|
|
|
if not success: |
|
106
|
|
|
check_output = pre_check_output |
|
107
|
|
|
for try_num in range(0, max_tries): |
|
108
|
|
|
if try_num > 0: |
|
109
|
|
|
console.print("Retrying") |
|
110
|
|
|
prompt = type_prompt(function_text, needs_typing, previous_error=check_output) |
|
111
|
|
|
answer = get_response_with_progress(INSTRUCTIONS, [], prompt) |
|
112
|
|
|
hints = parse_type_hint_answer(answer) |
|
113
|
|
|
|
|
114
|
|
|
insertions_for_function = add_type_hints(function_ast, hints, imports) |
|
115
|
|
|
if insertions_for_function: |
|
116
|
|
|
console.print(f"[cyan]Bot[/cyan]: {hints}") |
|
117
|
|
|
source_file.update_file(insertions_for_function, suffix=".shadow") |
|
118
|
|
|
(success, check_output) = run_check(check_command) |
|
119
|
|
|
if success: |
|
120
|
|
|
console.print("[green]Type checker passed[/green], keeping") |
|
121
|
|
|
return insertions_for_function |
|
122
|
|
|
else: |
|
123
|
|
|
console.out(check_output) |
|
124
|
|
|
console.print("\n[red]Type checker failed[/red], discarding") |
|
125
|
|
|
else: |
|
126
|
|
|
console.print("[cyan]Bot[/cyan]: No changes") |
|
127
|
|
|
# No retry if it didn't try to hint anything. |
|
128
|
|
|
return [] |
|
129
|
|
|
return [] |
|
130
|
|
|
|
|
131
|
|
|
|
|
132
|
|
|
@cli.command("type") |
|
133
|
|
|
@click.argument("file") |
|
134
|
|
|
def type_command(file): |
|
135
|
|
|
"""Insert type hints (Python only)""" |
|
136
|
|
|
from menderbot.typing import process_untyped_functions # Lazy import |
|
137
|
|
|
|
|
138
|
|
|
check_llm_consent() |
|
139
|
|
|
console.print("Running type-checker baseline") |
|
140
|
|
|
mypy_cmd = f"mypy --ignore-missing-imports --no-error-summary --soft-error-limit 10 '{file}'" |
|
141
|
|
|
(success, check_output) = run_check(f"{mypy_cmd}") |
|
142
|
|
|
if not success: |
|
143
|
|
|
console.print(check_output) |
|
144
|
|
|
console.print("Baseline failed, aborting.") |
|
145
|
|
|
return |
|
146
|
|
|
source_file = SourceFile(file) |
|
147
|
|
|
insertions = [] |
|
148
|
|
|
for function_ast, needs_typing in process_untyped_functions(source_file): |
|
149
|
|
|
insertions += try_function_type_hints( |
|
150
|
|
|
mypy_cmd, source_file, function_ast, needs_typing |
|
151
|
|
|
) |
|
152
|
|
|
if not insertions: |
|
153
|
|
|
console.print(f"No changes for '{file}.") |
|
154
|
|
|
return |
|
155
|
|
|
if not Confirm.ask(f"Write '{file}'?"): |
|
156
|
|
|
console.print("Skipping.") |
|
157
|
|
|
return |
|
158
|
|
|
source_file.update_file(insertions, suffix="") |
|
159
|
|
|
console.print("Done.") |
|
160
|
|
|
|
|
161
|
|
|
|
|
162
|
|
|
def get_response_with_progress(instructions, history, question): |
|
163
|
|
|
with Progress(transient=True) as progress: |
|
164
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
165
|
|
|
answer = get_response(instructions, history, question) |
|
166
|
|
|
progress.update(task, completed=True) |
|
167
|
|
|
return answer |
|
168
|
|
|
|
|
169
|
|
|
|
|
170
|
|
|
def generate_doc(code, file_extension): |
|
171
|
|
|
if not file_extension == ".py": |
|
172
|
|
|
# Until more types are supported. |
|
173
|
|
|
return None |
|
174
|
|
|
question = f""" |
|
175
|
|
|
Write a short Python docstring for this code. |
|
176
|
|
|
Do not include Arg lists. |
|
177
|
|
|
Respond with docstring only, no code. |
|
178
|
|
|
CODE: |
|
179
|
|
|
{code} |
|
180
|
|
|
""" |
|
181
|
|
|
from menderbot.code import function_indent, reindent # Lazy import |
|
182
|
|
|
|
|
183
|
|
|
with Progress(transient=True) as progress: |
|
184
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
185
|
|
|
doc_text = get_response(INSTRUCTIONS, [], question) |
|
186
|
|
|
progress.update(task, completed=True) |
|
187
|
|
|
if '"""' in doc_text: |
|
188
|
|
|
doc_text = doc_text[0 : doc_text.rfind('"""') + 3] |
|
189
|
|
|
doc_text = doc_text[doc_text.find('"""') :] |
|
190
|
|
|
indent = function_indent(code) |
|
191
|
|
|
doc_text = reindent(doc_text, indent) |
|
192
|
|
|
return doc_text |
|
193
|
|
|
|
|
194
|
|
|
|
|
195
|
|
|
@cli.command() |
|
196
|
|
|
@click.argument("file") |
|
197
|
|
|
def doc(file): |
|
198
|
|
|
"""Generate function-level documentation for the existing code (Python only).""" |
|
199
|
|
|
from menderbot.doc import document_file # Lazy import |
|
200
|
|
|
|
|
201
|
|
|
check_llm_consent() |
|
202
|
|
|
source_file = SourceFile(file) |
|
203
|
|
|
insertions = document_file(source_file, generate_doc) |
|
204
|
|
|
if not insertions: |
|
205
|
|
|
console.print(f"No updates found for '{file}'.") |
|
206
|
|
|
return |
|
207
|
|
|
if not Confirm.ask(f"Write '{file}'?"): |
|
208
|
|
|
console.print("Skipping.") |
|
209
|
|
|
return |
|
210
|
|
|
source_file.update_file(insertions, suffix="") |
|
211
|
|
|
console.print("Done.") |
|
212
|
|
|
|
|
213
|
|
|
|
|
214
|
|
View Code Duplication |
@cli.command() |
|
|
|
|
|
|
215
|
|
|
def review(): |
|
216
|
|
|
"""Review a code block or changeset and provide feedback.""" |
|
217
|
|
|
check_llm_consent() |
|
218
|
|
|
console.print("Reading diff from STDIN...") |
|
219
|
|
|
diff_text = click.get_text_stream("stdin").read() |
|
220
|
|
|
new_question = code_review_prompt(diff_text) |
|
221
|
|
|
with Progress(transient=True) as progress: |
|
222
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
223
|
|
|
response_1 = get_response(INSTRUCTIONS, [], new_question) |
|
224
|
|
|
progress.update(task, completed=True) |
|
225
|
|
|
console.print(f"[cyan]Bot[/cyan]:\n{response_1}") |
|
226
|
|
|
|
|
227
|
|
|
|
|
228
|
|
|
@cli.command() |
|
229
|
|
|
def commit(): |
|
230
|
|
|
"""Generate an informative commit message based on a changeset.""" |
|
231
|
|
|
check_llm_consent() |
|
232
|
|
|
diff_text = git_diff_head(staged=True) |
|
233
|
|
|
if not diff_text.strip(): |
|
234
|
|
|
console.print("[yellow]No changes staged. Please stage some then try again.") |
|
235
|
|
|
return |
|
236
|
|
|
new_question = change_list_prompt(diff_text) |
|
237
|
|
|
with Progress(transient=True) as progress: |
|
238
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
239
|
|
|
response_1 = get_response(INSTRUCTIONS, [], new_question) |
|
240
|
|
|
progress.update(task, completed=True) |
|
241
|
|
|
console.print(f"[cyan]Bot (initial pass)[/cyan]:\n{response_1}") |
|
242
|
|
|
|
|
243
|
|
|
if not response_1.strip(): |
|
244
|
|
|
console.print("[red]Didn't get a response for change list summary. Try again?") |
|
245
|
|
|
return |
|
246
|
|
|
question_2 = commit_msg_prompt(response_1) |
|
247
|
|
|
with Progress(transient=True) as progress: |
|
248
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
249
|
|
|
response_2 = unwrap_codeblock(get_response(INSTRUCTIONS, [], question_2)) |
|
250
|
|
|
progress.update(task, completed=True) |
|
251
|
|
|
console.print(f"[cyan]Bot[/cyan]:\n{response_2}") |
|
252
|
|
|
if not response_2.strip(): |
|
253
|
|
|
console.print("[red]Didn't get a response for commit message. Try again?") |
|
254
|
|
|
return |
|
255
|
|
|
git_commit(response_2) |
|
256
|
|
|
|
|
257
|
|
|
|
|
258
|
|
View Code Duplication |
@cli.command() |
|
|
|
|
|
|
259
|
|
|
def diff(): |
|
260
|
|
|
""" |
|
261
|
|
|
Summarize the differences between two versions of a codebase. Takes diff from STDIN. |
|
262
|
|
|
|
|
263
|
|
|
Try: |
|
264
|
|
|
git diff HEAD | menderbot diff |
|
265
|
|
|
git diff main..HEAD | menderbot diff |
|
266
|
|
|
""" |
|
267
|
|
|
check_llm_consent() |
|
268
|
|
|
console.print("Reading diff from STDIN...") |
|
269
|
|
|
diff_text = click.get_text_stream("stdin").read() |
|
270
|
|
|
new_question = change_list_prompt(diff_text) |
|
271
|
|
|
with Progress(transient=True) as progress: |
|
272
|
|
|
task = progress.add_task("[green]Processing...", total=None) |
|
273
|
|
|
response_1 = get_response(INSTRUCTIONS, [], new_question) |
|
274
|
|
|
progress.update(task, completed=True) |
|
275
|
|
|
console.print(f"[cyan]Bot[/cyan]:\n{response_1}") |
|
276
|
|
|
|
|
277
|
|
|
|
|
278
|
|
|
@cli.command() |
|
279
|
|
|
def ingest(): |
|
280
|
|
|
"""Index files in current repo to be used with 'ask' and 'chat'.""" |
|
281
|
|
|
check_llm_consent() |
|
282
|
|
|
ingest_repo() |
|
283
|
|
|
|
|
284
|
|
|
|
|
285
|
|
|
def check_llm_consent(): |
|
286
|
|
|
if not has_llm_consent(): |
|
287
|
|
|
console.print( |
|
288
|
|
|
"[red]Error[/red]: This repo does not have consent recorded in .menderbot-config.yaml" |
|
289
|
|
|
) |
|
290
|
|
|
if not has_config(): |
|
291
|
|
|
create_default_config() |
|
292
|
|
|
raise Abort() |
|
293
|
|
|
|
|
294
|
|
|
|
|
295
|
|
|
@cli.command() |
|
296
|
|
|
def check(): |
|
297
|
|
|
"""Verify we have what we need to run.""" |
|
298
|
|
|
git_dir = git_show_top_level() |
|
299
|
|
|
failed = False |
|
300
|
|
|
|
|
301
|
|
|
def check_condition(condition, ok_msg, failed_msg): |
|
302
|
|
|
nonlocal failed |
|
|
|
|
|
|
303
|
|
|
if condition: |
|
304
|
|
|
console.print(f"[green]OK[/green] {ok_msg}") |
|
305
|
|
|
else: |
|
306
|
|
|
console.print(f"[red]Failed[/red] {failed_msg}") |
|
307
|
|
|
failed = True |
|
308
|
|
|
|
|
309
|
|
|
check_condition( |
|
310
|
|
|
git_dir, f"Git repo {git_dir}", "Not in repo directory or git not installed" |
|
311
|
|
|
) |
|
312
|
|
|
check_condition( |
|
313
|
|
|
has_key(), |
|
314
|
|
|
f"OpenAI API key found in {key_env_var()}", |
|
315
|
|
|
f"OpenAI API key not found in {key_env_var()}", |
|
316
|
|
|
) |
|
317
|
|
|
if not has_config(): |
|
318
|
|
|
create_default_config("No .menderbot-config.yaml file found, creating...") |
|
319
|
|
|
check_condition( |
|
320
|
|
|
has_llm_consent(), |
|
321
|
|
|
"LLM consent configured for this repo", |
|
322
|
|
|
"LLM consent not recorded in .menderbot-config.yaml for this repo, please edit it", |
|
323
|
|
|
) |
|
324
|
|
|
if failed: |
|
325
|
|
|
sys.exit(1) |
|
326
|
|
|
|
|
327
|
|
|
|
|
328
|
|
|
if __name__ == "__main__": |
|
329
|
|
|
# pylint: disable-next=no-value-for-parameter |
|
330
|
|
|
cli(obj={}) |
|
331
|
|
|
|