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