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