Format: text/x-python
⬇ Download File (2.5 KB)
   1  import asyncio
   2  
   3  try:
   4      from jupyter_ai_magics import BaseProvider
   5      from langchain_community.llms import FakeListLLM
   6  except ImportError:
   7  
   8      class BaseProvider:
   9          pass
  10  
  11      class FakeListLLM:
  12          pass
  13  
  14  
  15  FIBONACCI = """\
  16  def fib(n):
  17      if n < 2: return n
  18      return fib(n - 1) + fib(n - 2)
  19  """
  20  
  21  
  22  class FibonacciCompletionProvider(BaseProvider, FakeListLLM):  # type: ignore[misc, valid-type]
  23  
  24      id = "my_provider"
  25      name = "My Provider"
  26      model_id_key = "model"
  27      models = ["model_a"]
  28  
  29      def __init__(self, **kwargs):
  30          kwargs["responses"] = ["This fake response will not be used for completion"]
  31          kwargs["model_id"] = "model_a"
  32          super().__init__(**kwargs)
  33  
  34      async def generate_inline_completions(self, request):
  35          raise ValueError("IPython only supports streaming models.")
  36  
  37      async def stream_inline_completions(self, request):
  38          from jupyter_ai.completions.models import (
  39              InlineCompletionList,
  40              InlineCompletionReply,
  41          )
  42  
  43          assert request.number > 0
  44          token = f"t{request.number}s0"
  45          last_line = request.prefix.rstrip("\n").splitlines()[-1]
  46  
  47          if not FIBONACCI.startswith(last_line):
  48              return
  49  
  50          yield InlineCompletionReply(
  51              list=InlineCompletionList(
  52                  items=[
  53                      {"insertText": "", "isIncomplete": True, "token": token},
  54                  ]
  55              ),
  56              reply_to=request.number,
  57          )
  58  
  59          async for reply in self._stream(
  60              FIBONACCI[len(last_line) :],
  61              request.number,
  62              token,
  63          ):
  64              yield reply
  65  
  66      async def _stream(self, sentence, request_number, token, start_with=""):
  67          from jupyter_ai.completions.models import InlineCompletionStreamChunk
  68  
  69          suggestion = start_with
  70  
  71          for fragment in sentence.split(" "):
  72              await asyncio.sleep(0.05)
  73              if suggestion:
  74                  suggestion += " "
  75              suggestion += fragment
  76              yield InlineCompletionStreamChunk(
  77                  type="stream",
  78                  response={"insertText": suggestion, "token": token},
  79                  reply_to=request_number,
  80                  done=False,
  81              )
  82  
  83          # finally, send a message confirming that we are done
  84          yield InlineCompletionStreamChunk(
  85              type="stream",
  86              response={"insertText": suggestion, "token": token},
  87              reply_to=request_number,
  88              done=True,
  89          )
  90