-
Notifications
You must be signed in to change notification settings - Fork 34
/
suggestions.py
102 lines (81 loc) · 3.22 KB
/
suggestions.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
"""
This module creates a custom exception handler that will send the error to OpenAI's
ChatGPT in order to help debug the code. It will also display the error in the
notebook as usual.
"""
from traceback import TracebackException
from types import TracebackType
from typing import Type
from IPython import InteractiveShell, get_ipython
from genai.context import PastAssists, PastErrors
from genai.display import GenaiMarkdown, Stage, can_handle_display_updates
from genai.generate import generate_exception_suggestion
# this function will be called on exceptions in any cell
def custom_exc(
shell: "InteractiveShell",
etype: Type[BaseException],
evalue: BaseException,
tb: TracebackType,
tb_offset=None,
):
# still show the error within the notebook, don't just swallow it
shell.showtraceback((etype, evalue, tb), tb_offset=tb_offset)
# On interrupt, just let it be
if etype == KeyboardInterrupt:
return
# On exit, release the user
elif etype == SystemExit:
return
try:
code = None
execution_count = shell.execution_count
In = shell.user_ns["In"]
history_manager = shell.history_manager
# If the history is available, use that as it has the raw inputs (including magics)
if (history_manager is not None) and (
execution_count == len(history_manager.input_hist_raw) - 1
):
code = shell.history_manager.input_hist_raw[execution_count]
# Fallback on In
elif In is not None and execution_count == len(In) - 1:
# Otherwise, use the current input buffer
code = In[execution_count]
# Otherwise history may not have been stored (store_history=False), so we should not send the
# code to GPT.
else:
code = None
gm = GenaiMarkdown(
"Let's see how we can fix this... 🔧",
stage=Stage.STARTING,
)
gm.display()
# Highly colorized tracebacks do not help GPT as much as a clean plaintext traceback.
formatted = TracebackException(etype, evalue, tb, limit=3).format(chain=True)
plaintext_traceback = "\n".join(formatted)
# Track context for future suggestions
PastErrors.add(execution_count, etype, evalue, tb)
PastAssists.add(execution_count, gm)
stream = can_handle_display_updates()
suggestion = generate_exception_suggestion(
code=code,
etype=etype,
evalue=evalue,
plaintext_traceback=plaintext_traceback,
stream=stream,
)
gm.stage = Stage.GENERATING
gm.message = "## 💡 Suggestion\n"
gm.consume(suggestion)
gm.stage = Stage.FINISHED
except Exception as e:
print("Error while trying to provide a suggestion: ", e)
except KeyboardInterrupt:
# If we have our heading, replace it with empty text and take out any stage information
if "gm" in locals():
gm.message = " "
gm.stage = None
def register(ipython=None):
"""Register the exception handler with the given IPython instance."""
if ipython is None:
ipython = get_ipython()
ipython.set_custom_exc((Exception,), custom_exc)