feat: :sparkles: support stablecoder with replicate LLM

4 files changed, 94 insertions, 4 deletions

diff --git a/continuedev/poetry.lock b/continuedev/poetry.lock
index 3754d121..d3140756 100644
--- a/continuedev/poetry.lock
+++ b/continuedev/poetry.lock
@@ -907,13 +907,13 @@ files = [
 
 [[package]]
 name = "openai"
-version = "0.27.6"
+version = "0.27.8"
 description = "Python client library for the OpenAI API"
 optional = false
 python-versions = ">=3.7.1"
 files = [
-    {file = "openai-0.27.6-py3-none-any.whl", hash = "sha256:1f07ed06f1cfc6c25126107193726fe4cf476edcc4e1485cd9eb708f068f2606"},
-    {file = "openai-0.27.6.tar.gz", hash = "sha256:63ca9f6ac619daef8c1ddec6d987fe6aa1c87a9bfdce31ff253204d077222375"},
+    {file = "openai-0.27.8-py3-none-any.whl", hash = "sha256:e0a7c2f7da26bdbe5354b03c6d4b82a2f34bd4458c7a17ae1a7092c3e397e03c"},
+    {file = "openai-0.27.8.tar.gz", hash = "sha256:2483095c7db1eee274cebac79e315a986c4e55207bb4fa7b82d185b3a2ed9536"},
 ]
 
 [package.dependencies]
@@ -928,6 +928,17 @@ embeddings = ["matplotlib", "numpy", "openpyxl (>=3.0.7)", "pandas (>=1.2.3)", "
 wandb = ["numpy", "openpyxl (>=3.0.7)", "pandas (>=1.2.3)", "pandas-stubs (>=1.1.0.11)", "wandb"]
 
 [[package]]
+name = "packaging"
+version = "23.1"
+description = "Core utilities for Python packages"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "packaging-23.1-py3-none-any.whl", hash = "sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61"},
+    {file = "packaging-23.1.tar.gz", hash = "sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f"},
+]
+
+[[package]]
 name = "pkgutil-resolve-name"
 version = "1.3.10"
 description = "Resolve a name to an object."
@@ -1275,6 +1286,25 @@ files = [
 ]
 
 [[package]]
+name = "replicate"
+version = "0.11.0"
+description = "Python client for Replicate"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "replicate-0.11.0-py3-none-any.whl", hash = "sha256:fbb8815068864dc822cd4fa7b6103d6f4089d6ef122abd6c3441ca0f0f110c46"},
+    {file = "replicate-0.11.0.tar.gz", hash = "sha256:4d54b5838c1552a6f76cc37c3af8d9a7998105382082d672acad31636ad443b5"},
+]
+
+[package.dependencies]
+packaging = "*"
+pydantic = ">1"
+requests = ">2"
+
+[package.extras]
+dev = ["black", "mypy", "pytest", "responses", "ruff"]
+
+[[package]]
 name = "requests"
 version = "2.29.0"
 description = "Python HTTP for Humans."
@@ -1819,4 +1849,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8.1"
-content-hash = "19ea56d05ab21d19f2fee6f837f7c8915ca1c597c392a539f43010dd0a4f6a78"
+content-hash = "5500ea86b06a96f5fe45939500936911e622043a67a3a5c3d02473463ff2fd6c"
diff --git a/continuedev/pyproject.toml b/continuedev/pyproject.toml
index e328b6cd..49b3c5ed 100644
--- a/continuedev/pyproject.toml
+++ b/continuedev/pyproject.toml
@@ -30,6 +30,7 @@ meilisearch-python-async = "^1.4.8"
 socksio = "^1.0.0"
 ripgrepy = "^2.0.0"
 bs4 = "^0.0.1"
+replicate = "^0.11.0"
 
 [tool.poetry.scripts]
 typegen = "src.continuedev.models.generate_json_schema:main" 
diff --git a/continuedev/src/continuedev/libs/llm/__init__.py b/continuedev/src/continuedev/libs/llm/__init__.py
index 40edb99b..70c67856 100644
--- a/continuedev/src/continuedev/libs/llm/__init__.py
+++ b/continuedev/src/continuedev/libs/llm/__init__.py
@@ -12,6 +12,9 @@ class LLM(ContinueBaseModel, ABC):
 
     system_message: Optional[str] = None
 
+    class Config:
+        arbitrary_types_allowed = True
+
     @abstractproperty
     def name(self):
         """Return the name of the LLM."""
diff --git a/continuedev/src/continuedev/libs/llm/replicate.py b/continuedev/src/continuedev/libs/llm/replicate.py
new file mode 100644
index 00000000..b13e2dec
--- /dev/null
+++ b/continuedev/src/continuedev/libs/llm/replicate.py
@@ -0,0 +1,56 @@
+from abc import abstractproperty
+from typing import List, Optional
+import replicate
+import concurrent.futures
+
+from ..util.count_tokens import DEFAULT_ARGS, count_tokens
+from ...core.main import ChatMessage
+from . import LLM
+
+
+class ReplicateLLM(LLM):
+    api_key: str
+    model: str = "nateraw/stablecode-completion-alpha-3b-4k:e82ebe958f0a5be6846d1a82041925767edb1d1f162596c643e48fbea332b1bb"
+    max_context_length: int = 2048
+
+    _client: replicate.Client = None
+
+    @property
+    def name(self):
+        return self.model
+
+    @property
+    def context_length(self):
+        return self.max_context_length
+
+    @property
+    def default_args(self):
+        return {**DEFAULT_ARGS, "model": self.name, "max_tokens": 1024}
+
+    def count_tokens(self, text: str):
+        return count_tokens(self.name, text)
+
+    async def start(self):
+        self._client = replicate.Client(api_token=self.api_key)
+
+    async def stop(self):
+        pass
+
+    async def complete(self, prompt: str, with_history: List[ChatMessage] = None, **kwargs):
+        output = self._client.run(self.model, input={"message": prompt})
+        completion = ''
+        for item in output:
+            completion += item
+
+        return completion
+
+    async def stream_complete(self, prompt, with_history: List[ChatMessage] = None, **kwargs):
+        for item in self._client.run(self.model, input={"message": prompt}):
+            yield item
+
+    async def stream_chat(self, messages: List[ChatMessage] = None, **kwargs):
+        for item in self._client.run(self.model, input={"message": messages[-1].content}):
+            yield {
+                "content": item,
+                "role": "assistant"
+            }