copying from old repo

6 files changed, 412 insertions, 0 deletions

diff --git a/continuedev/src/continuedev/libs/llm/__init__.py b/continuedev/src/continuedev/libs/llm/__init__.py
new file mode 100644
index 00000000..6bae2222
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/__init__.py
@@ -0,0 +1,22 @@
+from typing import Union
+from ...models.main import AbstractModel
+from pydantic import BaseModel
+
+
+class LLM(BaseModel):
+    system_message: Union[str, None] = None
+
+    def complete(self, prompt: str, **kwargs):
+        """Return the completion of the text with the given temperature."""
+        raise
+
+    def __call__(self, prompt: str, **kwargs):
+        return self.complete(prompt, **kwargs)
+
+    def fine_tune(self, pairs: list):
+        """Fine tune the model on the given prompt/completion pairs."""
+        raise NotImplementedError
+
+    def with_system_message(self, system_message: Union[str, None]):
+        """Return a new model with the given system message."""
+        raise NotImplementedError
diff --git a/continuedev/src/continuedev/libs/llm/hugging_face.py b/continuedev/src/continuedev/libs/llm/hugging_face.py
new file mode 100644
index 00000000..868cb560
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/hugging_face.py
@@ -0,0 +1,14 @@
+from .llm import LLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+class HuggingFace(LLM):
+    def __init__(self, model_path: str = "Salesforce/codegen-2B-mono"):
+        self.model_path = model_path
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path)
+    
+    def complete(self, prompt: str, **kwargs):
+        args = { "max_tokens": 100 } | kwargs
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        generated_ids = self.model.generate(input_ids, max_length=args["max_tokens"])
+        return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
\ No newline at end of file
diff --git a/continuedev/src/continuedev/libs/llm/openai.py b/continuedev/src/continuedev/libs/llm/openai.py
new file mode 100644
index 00000000..bb745e75
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/openai.py
@@ -0,0 +1,159 @@
+import asyncio
+import time
+from typing import Any, Dict, Generator, List, Union
+import openai
+import aiohttp
+from ..llm import LLM
+from pydantic import BaseModel, validator
+
+
+class OpenAI(LLM):
+    api_key: str
+    completion_count: int = 0
+    default_model: str = "text-davinci-003"
+
+    @validator("api_key", pre=True, always=True)
+    def validate_api_key(cls, v):
+        openai.api_key = v
+        return v
+
+    def with_system_message(self, system_message: Union[str, None]):
+        return OpenAI(api_key=self.api_key, system_message=system_message)
+
+    def stream_chat(self, messages, **kwargs) -> Generator[Union[Any, List, Dict], None, None]:
+        self.completion_count += 1
+        args = {"max_tokens": 512, "temperature": 0.5, "top_p": 1,
+                "frequency_penalty": 0, "presence_penalty": 0} | kwargs
+        args["stream"] = True
+        args["model"] = "gpt-3.5-turbo"
+
+        for chunk in openai.ChatCompletion.create(
+            messages=messages,
+            **args,
+        ):
+            if "content" in chunk.choices[0].delta:
+                yield chunk.choices[0].delta.content
+            else:
+                continue
+
+    def stream_complete(self, prompt: str, **kwargs) -> Generator[Union[Any, List, Dict], None, None]:
+        self.completion_count += 1
+        args = {"model": self.default_model, "max_tokens": 512, "temperature": 0.5,
+                "top_p": 1, "frequency_penalty": 0, "presence_penalty": 0, "suffix": None} | kwargs
+        args["stream"] = True
+
+        if args["model"] == "gpt-3.5-turbo":
+            generator = openai.ChatCompletion.create(
+                messages=[{
+                    "role": "user",
+                    "content": prompt
+                }],
+                **args,
+            )
+            for chunk in generator:
+                yield chunk.choices[0].message.content
+        else:
+            generator = openai.Completion.create(
+                prompt=prompt,
+                **args,
+            )
+            for chunk in generator:
+                yield chunk.choices[0].text
+
+    def complete(self, prompt: str, **kwargs) -> str:
+        t1 = time.time()
+
+        self.completion_count += 1
+        args = {"model": self.default_model, "max_tokens": 512, "temperature": 0.5, "top_p": 1,
+                "frequency_penalty": 0, "presence_penalty": 0, "stream": False} | kwargs
+
+        if args["model"] == "gpt-3.5-turbo":
+            messages = [{
+                "role": "user",
+                "content": prompt
+            }]
+            if self.system_message:
+                messages.insert(0, {
+                    "role": "system",
+                    "content": self.system_message
+                })
+            resp = openai.ChatCompletion.create(
+                messages=messages,
+                **args,
+            ).choices[0].message.content
+        else:
+            resp = openai.Completion.create(
+                prompt=prompt,
+                **args,
+            ).choices[0].text
+
+        t2 = time.time()
+        print("Completion time:", t2 - t1)
+        return resp
+
+    def edit(self, inp: str, instruction: str) -> str:
+        try:
+            resp = openai.Edit.create(
+                input=inp,
+                instruction=instruction,
+                model='text-davinci-edit-001'
+            ).choices[0].text
+            return resp
+        except Exception as e:
+            print("OpenAI error:", e)
+            raise e
+
+    def parallel_edit(self, inputs: list[str], instructions: Union[List[str], str], **kwargs) -> list[str]:
+        args = {"temperature": 0.5, "top_p": 1} | kwargs
+        args['model'] = 'text-davinci-edit-001'
+
+        async def fn():
+            async with aiohttp.ClientSession() as session:
+                tasks = []
+
+                async def get(input, instruction):
+                    async with session.post("https://api.openai.com/v1/edits", headers={
+                        "Content-Type": "application/json",
+                        "Authorization": "Bearer " + self.api_key
+                    }, json={"model": args["model"], "input": input, "instruction": instruction, "temperature": args["temperature"], "max_tokens": args["max_tokens"], "suffix": args["suffix"]}) as resp:
+                        json = await resp.json()
+                        if "error" in json:
+                            print("ERROR IN GPT-3 RESPONSE: ", json)
+                            return None
+                        return json["choices"][0]["text"]
+
+                for i in range(len(inputs)):
+                    tasks.append(get(inputs[i], instructions[i] if isinstance(
+                        instructions, list) else instructions))
+
+                return await asyncio.gather(*tasks)
+
+        return asyncio.run(fn())
+
+    def parallel_complete(self, prompts: list[str], suffixes: Union[list[str], None] = None, **kwargs) -> list[str]:
+        self.completion_count += len(prompts)
+        args = {"model": self.default_model, "max_tokens": 512, "temperature": 0.5,
+                "top_p": 1, "frequency_penalty": 0, "presence_penalty": 0} | kwargs
+
+        async def fn():
+            async with aiohttp.ClientSession() as session:
+                tasks = []
+
+                async def get(prompt, suffix):
+                    async with session.post("https://api.openai.com/v1/completions", headers={
+                        "Content-Type": "application/json",
+                        "Authorization": "Bearer " + self.api_key
+                    }, json={"model": args["model"], "prompt": prompt, "temperature": args["temperature"], "max_tokens": args["max_tokens"], "suffix": suffix}) as resp:
+                        json = await resp.json()
+                        if "error" in json:
+                            print("ERROR IN GPT-3 RESPONSE: ", json)
+                            return None
+                        return json["choices"][0]["text"]
+
+                for i in range(len(prompts)):
+                    tasks.append(asyncio.ensure_future(
+                        get(prompts[i], suffixes[i] if suffixes else None)))
+
+                return await asyncio.gather(*tasks)
+
+        return asyncio.run(fn())
diff --git a/continuedev/src/continuedev/libs/llm/prompt_utils.py b/continuedev/src/continuedev/libs/llm/prompt_utils.py
new file mode 100644
index 00000000..51b64732
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/prompt_utils.py
@@ -0,0 +1,71 @@
+from typing import Dict, List, Union
+from ...models.filesystem import RangeInFileWithContents
+from ...models.filesystem_edit import FileEdit
+
+
+class MarkdownStyleEncoderDecoder:
+    # Filename -> the part of the file you care about
+    range_in_files: List[RangeInFileWithContents]
+
+    def __init__(self, range_in_files: List[RangeInFileWithContents]):
+        self.range_in_files = range_in_files
+
+    def encode(self) -> str:
+        return "\n\n".join([
+            f"File ({rif.filepath})\n```\n{rif.contents}\n```"
+            for rif in self.range_in_files
+        ])
+
+    def _suggestions_to_file_edits(self, suggestions: Dict[str, str]) -> List[FileEdit]:
+        file_edits: List[FileEdit] = []
+        for suggestion_filepath, suggestion in suggestions.items():
+            matching_rifs = list(
+                filter(lambda r: r.filepath == suggestion_filepath, self.range_in_files))
+            if (len(matching_rifs) > 0):
+                range_in_file = matching_rifs[0]
+                file_edits.append(FileEdit(
+                    range=range_in_file.range,
+                    filepath=range_in_file.filepath,
+                    replacement=suggestion
+                ))
+
+        return file_edits
+
+    def _decode_to_suggestions(self, completion: str) -> Dict[str, str]:
+        if len(self.range_in_files) == 0:
+            return {}
+
+        if not '```' in completion:
+            completion = "```\n" + completion + "\n```"
+        if completion.strip().splitlines()[0].strip() == '```':
+            first_filepath = self.range_in_files[0].filepath
+            completion = f"File ({first_filepath})\n" + completion
+
+        suggestions: Dict[str, str] = {}
+        current_file_lines: List[str] = []
+        current_filepath: Union[str, None] = None
+        last_was_file = False
+        inside_file = False
+        for line in completion.splitlines():
+            if line.strip().startswith("File ("):
+                last_was_file = True
+                current_filepath = line.strip()[6:-1]
+            elif last_was_file and line.startswith("```"):
+                last_was_file = False
+                inside_file = True
+            elif inside_file:
+                if line.startswith("```"):
+                    inside_file = False
+                    suggestions[current_filepath] = "\n".join(
+                        current_file_lines)
+                    current_file_lines = []
+                    current_filepath = None
+                else:
+                    current_file_lines.append(line)
+
+        return suggestions
+
+    def decode(self, completion: str) -> List[FileEdit]:
+        suggestions = self._decode_to_suggestions(completion)
+        file_edits = self._suggestions_to_file_edits(suggestions)
+        return file_edits
diff --git a/continuedev/src/continuedev/libs/llm/prompters.py b/continuedev/src/continuedev/libs/llm/prompters.py
new file mode 100644
index 00000000..04e9885a
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/prompters.py
@@ -0,0 +1,112 @@
+from typing import Any, Callable, List, Tuple, Union
+from ..llm import LLM
+from .openai import OpenAI
+
+
+def cls_method_to_str(cls_name: str, init: str, method: str) -> str:
+    """Convert class and method info to formatted code"""
+    return f"""class {cls_name}:
+{init}
+{method}"""
+
+
+# Prompter classes
+class Prompter:
+    def __init__(self, llm: LLM = None):
+        if llm is None:
+            self.llm = OpenAI()
+        else:
+            self.llm = llm
+
+    def _compile_prompt(self, inp: Any) -> Tuple[str, str, Union[str, None]]:
+        "Takes input and returns prompt, prefix, suffix"
+        raise NotImplementedError
+
+    def complete(self, inp: Any, **kwargs) -> str:
+        prompt, prefix, suffix = self._compile_prompt(inp)
+        resp = self.llm.complete(prompt + prefix, suffix=suffix, **kwargs)
+        return prefix + resp + (suffix or "")
+
+    def __call__(self, inp: Any, **kwargs) -> str:
+        return self.complete(inp, **kwargs)
+
+    def parallel_complete(self, inps: List[Any]) -> List[str]:
+        prompts = []
+        prefixes = []
+        suffixes = []
+        for inp in inps:
+            prompt, prefix, suffix = self._compile_prompt(inp)
+            prompts.append(prompt)
+            prefixes.append(prefix)
+            suffixes.append(suffix)
+
+        resps = self.llm.parallel_complete(
+            [prompt + prefix for prompt, prefix in zip(prompts, prefixes)], suffixes=suffixes)
+        return [prefix + resp + (suffix or "") for prefix, resp, suffix in zip(prefixes, resps, suffixes)]
+
+
+class MixedPrompter(Prompter):
+    def __init__(self, prompters: List[Prompter], router: Callable[[Any], int], llm: LLM = None):
+        super().__init__(llm=llm)
+        self.prompters = prompters
+        self.router = router
+
+    def _compile_prompt(self, inp: Any) -> Tuple[str, str, Union[str, None]]:
+        prompter = self.prompters[self.router(inp)]
+        return prompter._compile_prompt(inp)
+
+    def complete(self, inp: Any, **kwargs) -> str:
+        prompter = self.prompters[self.router(inp)]
+        return prompter.complete(inp, **kwargs)
+
+
+class SimplePrompter(Prompter):
+    def __init__(self, prompt_fn: Callable[[Any], str], llm: LLM = None):
+        super().__init__(llm=llm)
+        self.prompt_fn = prompt_fn
+
+    def _compile_prompt(self, inp: Any) -> Tuple[str, str, Union[str, None]]:
+        return self.prompt_fn(inp), "", None
+
+
+class FormatStringPrompter(SimplePrompter):
+    """Pass a formatted string, and the input should be a dict with the keys to format"""
+
+    def __init__(self, prompt: str, llm: LLM = None):
+        super().__init__(lambda inp: prompt.format(**inp), llm=llm)
+
+
+class BasicCommentPrompter(SimplePrompter):
+    def __init__(self, comment: str, llm: LLM = None):
+        super().__init__(lambda inp: f"""{inp}
+
+# {comment}""", llm=llm)
+
+
+class EditPrompter(Prompter):
+    def __init__(self, prompt_fn: Callable[[Any], Tuple[str, str]], llm: LLM = None):
+        super().__init__(llm=llm)
+        self.prompt_fn = prompt_fn
+
+    def complete(self, inp: str, **kwargs) -> str:
+        inp, instruction = self.prompt_fn(inp)
+        return self.llm.edit(inp, instruction, **kwargs)
+
+    def parallel_complete(self, inps: List[Any]) -> List[str]:
+        prompts = []
+        instructions = []
+        for inp in inps:
+            prompt, instruction = self.prompt_fn(inp)
+            prompts.append(prompt)
+            instructions.append(instruction)
+
+        return self.llm.parallel_edit(prompts, instructions)
+
+
+class InsertPrompter(Prompter):
+    def __init__(self, prompt_fn: Callable[[Any], Tuple[str, str, str]], llm: LLM = None):
+        super().__init__(llm=llm)
+        self.prompt_fn = prompt_fn
+
+    def _compile_prompt(self, inp: Any) -> Tuple[str, str, Union[str, None]]:
+        return self.prompt_fn(inp)
diff --git a/continuedev/src/continuedev/libs/llm/utils.py b/continuedev/src/continuedev/libs/llm/utils.py
new file mode 100644
index 00000000..76240d4e
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/utils.py
@@ -0,0 +1,34 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import GPT2TokenizerFast
+
+gpt2_tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
+def count_tokens(text: str) -> int:
+    return len(gpt2_tokenizer.encode(text))
+
+prices = {
+    # All prices are per 1k tokens
+    "fine-tune-train": {
+        "davinci": 0.03,
+        "curie": 0.03,
+        "babbage": 0.0006,
+        "ada": 0.0004,
+    },
+    "completion": {
+        "davinci": 0.02,
+        "curie": 0.002,
+        "babbage": 0.0005,
+        "ada": 0.0004,
+    },
+    "fine-tune-completion": {
+        "davinci": 0.12,
+        "curie": 0.012,
+        "babbage": 0.0024,
+        "ada": 0.0016,
+    },
+    "embedding": {
+        "ada": 0.0004
+    }
+}
+
+def get_price(text: str, model: str="davinci", task: str="completion") -> float:
+    return count_tokens(text) * prices[task][model] / 1000
\ No newline at end of file