Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import re
2import asyncio
3import tokenize
4from io import StringIO
5from typing import ClassVar, Any
6from collections.abc import Callable, Generator
7import warnings
9import prompt_toolkit
10from prompt_toolkit.buffer import Buffer
11from prompt_toolkit.key_binding import KeyPressEvent
12from prompt_toolkit.key_binding.bindings import named_commands as nc
13from prompt_toolkit.auto_suggest import AutoSuggestFromHistory, Suggestion
14from prompt_toolkit.document import Document
15from prompt_toolkit.history import History
16from prompt_toolkit.shortcuts import PromptSession
17from prompt_toolkit.layout.processors import (
18 Processor,
19 Transformation,
20 TransformationInput,
21)
23from IPython.core.getipython import get_ipython
24from IPython.utils.tokenutil import generate_tokens
26from .filters import pass_through
29def _get_query(document: Document):
30 return document.lines[document.cursor_position_row]
33class AppendAutoSuggestionInAnyLine(Processor):
34 """
35 Append the auto suggestion to lines other than the last (appending to the
36 last line is natively supported by the prompt toolkit).
38 This has a private `_debug` attribute that can be set to True to display
39 debug information as virtual suggestion on the end of any line. You can do
40 so with:
42 >>> from IPython.terminal.shortcuts.auto_suggest import AppendAutoSuggestionInAnyLine
43 >>> AppendAutoSuggestionInAnyLine._debug = True
45 """
47 _debug: ClassVar[bool] = False
49 def __init__(self, style: str = "class:auto-suggestion") -> None:
50 self.style = style
52 def apply_transformation(self, ti: TransformationInput) -> Transformation:
53 """
54 Apply transformation to the line that is currently being edited.
56 This is a variation of the original implementation in prompt toolkit
57 that allows to not only append suggestions to any line, but also to show
58 multi-line suggestions.
60 As transformation are applied on a line-by-line basis; we need to trick
61 a bit, and elide any line that is after the line we are currently
62 editing, until we run out of completions. We cannot shift the existing
63 lines
65 There are multiple cases to handle:
67 The completions ends before the end of the buffer:
68 We can resume showing the normal line, and say that some code may
69 be hidden.
71 The completions ends at the end of the buffer
72 We can just say that some code may be hidden.
74 And separately:
76 The completions ends beyond the end of the buffer
77 We need to both say that some code may be hidden, and that some
78 lines are not shown.
80 """
81 last_line_number = ti.document.line_count - 1
82 is_last_line = ti.lineno == last_line_number
84 noop = lambda text: Transformation(
85 fragments=ti.fragments + [(self.style, " " + text if self._debug else "")]
86 )
87 if ti.document.line_count == 1:
88 return noop("noop:oneline")
89 if ti.document.cursor_position_row == last_line_number and is_last_line:
90 # prompt toolkit already appends something; just leave it be
91 return noop("noop:last line and cursor")
93 # first everything before the current line is unchanged.
94 if ti.lineno < ti.document.cursor_position_row:
95 return noop("noop:before cursor")
97 buffer = ti.buffer_control.buffer
98 if not buffer.suggestion or not ti.document.is_cursor_at_the_end_of_line:
99 return noop("noop:not eol")
101 delta = ti.lineno - ti.document.cursor_position_row
102 suggestions = buffer.suggestion.text.splitlines()
104 if len(suggestions) == 0:
105 return noop("noop: no suggestions")
107 if prompt_toolkit.VERSION < (3, 0, 49):
108 if len(suggestions) > 1 and prompt_toolkit.VERSION < (3, 0, 49):
109 if ti.lineno == ti.document.cursor_position_row:
110 return Transformation(
111 fragments=ti.fragments
112 + [
113 (
114 "red",
115 "(Cannot show multiline suggestion; requires prompt_toolkit > 3.0.49)",
116 )
117 ]
118 )
119 else:
120 return Transformation(fragments=ti.fragments)
121 elif len(suggestions) == 1:
122 if ti.lineno == ti.document.cursor_position_row:
123 return Transformation(
124 fragments=ti.fragments + [(self.style, suggestions[0])]
125 )
126 return Transformation(fragments=ti.fragments)
128 if delta == 0:
129 suggestion = suggestions[0]
130 return Transformation(fragments=ti.fragments + [(self.style, suggestion)])
131 if is_last_line:
132 if delta < len(suggestions):
133 suggestion = f"… rest of suggestion ({len(suggestions) - delta} lines) and code hidden"
134 return Transformation([(self.style, suggestion)])
136 n_elided = len(suggestions)
137 for i in range(len(suggestions)):
138 ll = ti.get_line(last_line_number - i)
139 el = "".join(l[1] for l in ll).strip()
140 if el:
141 break
142 else:
143 n_elided -= 1
144 if n_elided:
145 return Transformation([(self.style, f"… {n_elided} line(s) hidden")])
146 else:
147 return Transformation(
148 ti.get_line(last_line_number - len(suggestions) + 1)
149 + ([(self.style, "shift-last-line")] if self._debug else [])
150 )
152 elif delta < len(suggestions):
153 suggestion = suggestions[delta]
154 return Transformation([(self.style, suggestion)])
155 else:
156 shift = ti.lineno - len(suggestions) + 1
157 return Transformation(ti.get_line(shift))
160class NavigableAutoSuggestFromHistory(AutoSuggestFromHistory):
161 """
162 A subclass of AutoSuggestFromHistory that allow navigation to next/previous
163 suggestion from history. To do so it remembers the current position, but it
164 state need to carefully be cleared on the right events.
165 """
167 skip_lines: int
168 _connected_apps: list[PromptSession]
170 # handle to the currently running llm task that appends suggestions to the
171 # current buffer; we keep a handle to it in order to cancel it when there is a cursor movement, or
172 # another request.
173 _llm_task: asyncio.Task | None = None
175 # This is the constructor of the LLM provider from jupyter-ai
176 # to which we forward the request to generate inline completions.
177 _init_llm_provider: Callable | None
179 _llm_provider_instance: Any | None
180 _llm_prefixer: Callable = lambda self, x: ""
182 def __init__(self):
183 super().__init__()
184 self.skip_lines = 0
185 self._connected_apps = []
186 self._llm_provider_instance = None
187 self._init_llm_provider = None
188 self._request_number = 0
190 def reset_history_position(self, _: Buffer) -> None:
191 self.skip_lines = 0
193 def disconnect(self) -> None:
194 self._cancel_running_llm_task()
195 for pt_app in self._connected_apps:
196 pt_app.default_buffer.on_text_insert.remove_handler(self.reset_history_position)
197 pt_app.default_buffer.on_cursor_position_changed.remove_handler(self._dismiss)
198 self._connected_apps = []
200 def connect(self, pt_app: PromptSession) -> None:
201 self._connected_apps.append(pt_app)
202 # note: `on_text_changed` could be used for a bit different behaviour
203 # on character deletion (i.e. resetting history position on backspace)
204 pt_app.default_buffer.on_text_insert.add_handler(self.reset_history_position)
205 pt_app.default_buffer.on_cursor_position_changed.add_handler(self._dismiss)
207 def get_suggestion(
208 self, buffer: Buffer, document: Document
209 ) -> Suggestion | None:
210 text = _get_query(document)
212 if text.strip():
213 for suggestion, _ in self._find_next_match(
214 text, self.skip_lines, buffer.history
215 ):
216 return Suggestion(suggestion)
218 return None
220 def _dismiss(self, buffer, *args, **kwargs) -> None:
221 self._cancel_running_llm_task()
222 buffer.suggestion = None
224 def _find_match(
225 self, text: str, skip_lines: float, history: History, previous: bool
226 ) -> Generator[tuple[str, float], None, None]:
227 """
228 text : str
229 Text content to find a match for, the user cursor is most of the
230 time at the end of this text.
231 skip_lines : float
232 number of items to skip in the search, this is used to indicate how
233 far in the list the user has navigated by pressing up or down.
234 The float type is used as the base value is +inf
235 history : History
236 prompt_toolkit History instance to fetch previous entries from.
237 previous : bool
238 Direction of the search, whether we are looking previous match
239 (True), or next match (False).
241 Yields
242 ------
243 Tuple with:
244 str:
245 current suggestion.
246 float:
247 will actually yield only ints, which is passed back via skip_lines,
248 which may be a +inf (float)
251 """
252 line_number = -1
253 for string in reversed(list(history.get_strings())):
254 for line in reversed(string.splitlines()):
255 line_number += 1
256 if not previous and line_number < skip_lines:
257 continue
258 # do not return empty suggestions as these
259 # close the auto-suggestion overlay (and are useless)
260 if line.startswith(text) and len(line) > len(text):
261 yield line[len(text) :], line_number
262 if previous and line_number >= skip_lines:
263 return
265 def _find_next_match(
266 self, text: str, skip_lines: float, history: History
267 ) -> Generator[tuple[str, float], None, None]:
268 return self._find_match(text, skip_lines, history, previous=False)
270 def _find_previous_match(self, text: str, skip_lines: float, history: History):
271 return reversed(
272 list(self._find_match(text, skip_lines, history, previous=True))
273 )
275 def up(self, query: str, other_than: str, history: History) -> None:
276 self._cancel_running_llm_task()
277 for suggestion, line_number in self._find_next_match(
278 query, self.skip_lines, history
279 ):
280 # if user has history ['very.a', 'very', 'very.b'] and typed 'very'
281 # we want to switch from 'very.b' to 'very.a' because a) if the
282 # suggestion equals current text, prompt-toolkit aborts suggesting
283 # b) user likely would not be interested in 'very' anyways (they
284 # already typed it).
285 if query + suggestion != other_than:
286 self.skip_lines = line_number
287 break
288 else:
289 # no matches found, cycle back to beginning
290 self.skip_lines = 0
292 def down(self, query: str, other_than: str, history: History) -> None:
293 self._cancel_running_llm_task()
294 for suggestion, line_number in self._find_previous_match(
295 query, self.skip_lines, history
296 ):
297 if query + suggestion != other_than:
298 self.skip_lines = line_number
299 break
300 else:
301 # no matches found, cycle to end
302 for suggestion, line_number in self._find_previous_match(
303 query, float("Inf"), history
304 ):
305 if query + suggestion != other_than:
306 self.skip_lines = line_number
307 break
309 def _cancel_running_llm_task(self) -> None:
310 """
311 Try to cancel the currently running llm_task if exists, and set it to None.
312 """
313 if self._llm_task is not None:
314 if self._llm_task.done():
315 self._llm_task = None
316 return
317 cancelled = self._llm_task.cancel()
318 if cancelled:
319 self._llm_task = None
320 if not cancelled:
321 warnings.warn(
322 "LLM task not cancelled, does your provider support cancellation?"
323 )
325 @property
326 def _llm_provider(self):
327 """Lazy-initialized instance of the LLM provider.
329 Do not use in the constructor, as `_init_llm_provider` can trigger slow side-effects.
330 """
331 if self._llm_provider_instance is None and self._init_llm_provider:
332 self._llm_provider_instance = self._init_llm_provider()
333 return self._llm_provider_instance
335 async def _trigger_llm(self, buffer) -> None:
336 """
337 This will ask the current llm provider a suggestion for the current buffer.
339 If there is a currently running llm task, it will cancel it.
340 """
341 # we likely want to store the current cursor position, and cancel if the cursor has moved.
342 try:
343 import jupyter_ai_magics
344 except ModuleNotFoundError:
345 jupyter_ai_magics = None
346 if not self._llm_provider:
347 warnings.warn("No LLM provider found, cannot trigger LLM completions")
348 return
349 if jupyter_ai_magics is None:
350 warnings.warn("LLM Completion requires `jupyter_ai_magics` to be installed")
352 self._cancel_running_llm_task()
354 async def error_catcher(buffer):
355 """
356 This catches and log any errors, as otherwise this is just
357 lost in the void of the future running task.
358 """
359 try:
360 await self._trigger_llm_core(buffer)
361 except Exception as e:
362 get_ipython().log.error("error %s", e)
363 raise
365 # here we need a cancellable task so we can't just await the error caught
366 self._llm_task = asyncio.create_task(error_catcher(buffer))
367 try:
368 await self._llm_task
369 except (asyncio.CancelledError, Exception):
370 pass
372 async def _trigger_llm_core(self, buffer: Buffer):
373 """
374 This is the core of the current llm request.
376 Here we build a compatible `InlineCompletionRequest` and ask the llm
377 provider to stream it's response back to us iteratively setting it as
378 the suggestion on the current buffer.
380 Unlike with JupyterAi, as we do not have multiple cells, the cell id
381 is always set to `None`.
383 We set the prefix to the current cell content, but could also insert the
384 rest of the history or even just the non-fail history.
386 In the same way, we do not have cell id.
388 LLM provider may return multiple suggestion stream, but for the time
389 being we only support one.
391 Here we make the assumption that the provider will have
392 stream_inline_completions, I'm not sure it is the case for all
393 providers.
394 """
395 try:
396 import jupyter_ai.completions.models as jai_models
397 except ModuleNotFoundError:
398 jai_models = None
400 if not jai_models:
401 raise ValueError("jupyter-ai is not installed")
403 if not self._llm_provider:
404 raise ValueError("No LLM provider found, cannot trigger LLM completions")
406 hm = buffer.history.shell.history_manager
407 prefix = self._llm_prefixer(hm)
408 get_ipython().log.debug("prefix: %s", prefix)
410 self._request_number += 1
411 request_number = self._request_number
413 request = jai_models.InlineCompletionRequest(
414 number=request_number,
415 prefix=prefix + buffer.document.text_before_cursor,
416 suffix=buffer.document.text_after_cursor,
417 mime="text/x-python",
418 stream=True,
419 path=None,
420 language="python",
421 cell_id=None,
422 )
424 async for reply_and_chunks in self._llm_provider.stream_inline_completions(
425 request
426 ):
427 if self._request_number != request_number:
428 # If a new suggestion was requested, skip processing this one.
429 return
430 if isinstance(reply_and_chunks, jai_models.InlineCompletionReply):
431 if len(reply_and_chunks.list.items) > 1:
432 raise ValueError(
433 "Terminal IPython cannot deal with multiple LLM suggestions at once"
434 )
435 buffer.suggestion = Suggestion(
436 reply_and_chunks.list.items[0].insertText
437 )
438 buffer.on_suggestion_set.fire()
439 elif isinstance(reply_and_chunks, jai_models.InlineCompletionStreamChunk):
440 buffer.suggestion = Suggestion(reply_and_chunks.response.insertText)
441 buffer.on_suggestion_set.fire()
442 return
445async def llm_autosuggestion(event: KeyPressEvent):
446 """
447 Ask the AutoSuggester from history to delegate to ask an LLM for completion
449 This will first make sure that the current buffer have _MIN_LINES (7)
450 available lines to insert the LLM completion
452 Provisional as of 8.32, may change without warnings
454 """
455 _MIN_LINES = 5
456 provider = get_ipython().auto_suggest
457 if not isinstance(provider, NavigableAutoSuggestFromHistory):
458 return
459 doc = event.current_buffer.document
460 lines_to_insert = max(0, _MIN_LINES - doc.line_count + doc.cursor_position_row)
461 for _ in range(lines_to_insert):
462 event.current_buffer.insert_text("\n", move_cursor=False, fire_event=False)
464 await provider._trigger_llm(event.current_buffer)
467def accept_or_jump_to_end(event: KeyPressEvent):
468 """Apply autosuggestion or jump to end of line."""
469 buffer = event.current_buffer
470 d = buffer.document
471 after_cursor = d.text[d.cursor_position :]
472 lines = after_cursor.split("\n")
473 end_of_current_line = lines[0].strip()
474 suggestion = buffer.suggestion
475 if (suggestion is not None) and (suggestion.text) and (end_of_current_line == ""):
476 buffer.insert_text(suggestion.text)
477 else:
478 nc.end_of_line(event)
481def accept(event: KeyPressEvent):
482 """Accept autosuggestion"""
483 buffer = event.current_buffer
484 suggestion = buffer.suggestion
485 if suggestion:
486 buffer.insert_text(suggestion.text)
487 else:
488 nc.forward_char(event)
491def discard(event: KeyPressEvent):
492 """Discard autosuggestion"""
493 buffer = event.current_buffer
494 buffer.suggestion = None
497def accept_word(event: KeyPressEvent):
498 """Fill partial autosuggestion by word"""
499 buffer = event.current_buffer
500 suggestion = buffer.suggestion
501 if suggestion:
502 t = re.split(r"(\S+\s+)", suggestion.text)
503 buffer.insert_text(next((x for x in t if x), ""))
504 else:
505 nc.forward_word(event)
508def accept_character(event: KeyPressEvent):
509 """Fill partial autosuggestion by character"""
510 b = event.current_buffer
511 suggestion = b.suggestion
512 if suggestion and suggestion.text:
513 b.insert_text(suggestion.text[0])
516def accept_and_keep_cursor(event: KeyPressEvent):
517 """Accept autosuggestion and keep cursor in place"""
518 buffer = event.current_buffer
519 old_position = buffer.cursor_position
520 suggestion = buffer.suggestion
521 if suggestion:
522 buffer.insert_text(suggestion.text)
523 buffer.cursor_position = old_position
526def accept_and_move_cursor_left(event: KeyPressEvent):
527 """Accept autosuggestion and move cursor left in place"""
528 accept_and_keep_cursor(event)
529 nc.backward_char(event)
532def _update_hint(buffer: Buffer):
533 if buffer.auto_suggest:
534 suggestion = buffer.auto_suggest.get_suggestion(buffer, buffer.document)
535 buffer.suggestion = suggestion
538def backspace_and_resume_hint(event: KeyPressEvent):
539 """Resume autosuggestions after deleting last character"""
540 nc.backward_delete_char(event)
541 _update_hint(event.current_buffer)
544def resume_hinting(event: KeyPressEvent):
545 """Resume autosuggestions"""
546 pass_through.reply(event)
547 # Order matters: if update happened first and event reply second, the
548 # suggestion would be auto-accepted if both actions are bound to same key.
549 _update_hint(event.current_buffer)
552def up_and_update_hint(event: KeyPressEvent):
553 """Go up and update hint"""
554 current_buffer = event.current_buffer
556 current_buffer.auto_up(count=event.arg)
557 _update_hint(current_buffer)
560def down_and_update_hint(event: KeyPressEvent):
561 """Go down and update hint"""
562 current_buffer = event.current_buffer
564 current_buffer.auto_down(count=event.arg)
565 _update_hint(current_buffer)
568def accept_token(event: KeyPressEvent):
569 """Fill partial autosuggestion by token"""
570 b = event.current_buffer
571 suggestion = b.suggestion
573 if suggestion:
574 prefix = _get_query(b.document)
575 text = prefix + suggestion.text
577 tokens: list[str | None] = [None, None, None]
578 substrings = [""]
579 i = 0
581 for token in generate_tokens(StringIO(text).readline):
582 if token.type == tokenize.NEWLINE:
583 index = len(text)
584 else:
585 index = text.index(token[1], len(substrings[-1]))
586 substrings.append(text[:index])
587 tokenized_so_far = substrings[-1]
588 if tokenized_so_far.startswith(prefix):
589 if i == 0 and len(tokenized_so_far) > len(prefix):
590 tokens[0] = tokenized_so_far[len(prefix) :]
591 substrings.append(tokenized_so_far)
592 i += 1
593 tokens[i] = token[1]
594 if i == 2:
595 break
596 i += 1
598 if tokens[0]:
599 to_insert: str
600 insert_text = substrings[-2]
601 if tokens[1] and len(tokens[1]) == 1:
602 insert_text = substrings[-1]
603 to_insert = insert_text[len(prefix) :]
604 b.insert_text(to_insert)
605 return
607 nc.forward_word(event)
610Provider = AutoSuggestFromHistory | NavigableAutoSuggestFromHistory | None
613def _swap_autosuggestion(
614 buffer: Buffer,
615 provider: NavigableAutoSuggestFromHistory,
616 direction_method: Callable,
617):
618 """
619 We skip most recent history entry (in either direction) if it equals the
620 current autosuggestion because if user cycles when auto-suggestion is shown
621 they most likely want something else than what was suggested (otherwise
622 they would have accepted the suggestion).
623 """
624 suggestion = buffer.suggestion
625 if not suggestion:
626 return
628 query = _get_query(buffer.document)
629 current = query + suggestion.text
631 direction_method(query=query, other_than=current, history=buffer.history)
633 new_suggestion = provider.get_suggestion(buffer, buffer.document)
634 buffer.suggestion = new_suggestion
637def swap_autosuggestion_up(event: KeyPressEvent):
638 """Get next autosuggestion from history."""
639 shell = get_ipython()
640 provider = shell.auto_suggest
642 if not isinstance(provider, NavigableAutoSuggestFromHistory):
643 return
645 return _swap_autosuggestion(
646 buffer=event.current_buffer, provider=provider, direction_method=provider.up
647 )
650def swap_autosuggestion_down(event: KeyPressEvent):
651 """Get previous autosuggestion from history."""
652 shell = get_ipython()
653 provider = shell.auto_suggest
655 if not isinstance(provider, NavigableAutoSuggestFromHistory):
656 return
658 return _swap_autosuggestion(
659 buffer=event.current_buffer,
660 provider=provider,
661 direction_method=provider.down,
662 )