1 files changed, 125 insertions, 0 deletions
diff --git a/server/continuedev/libs/llm/together.py b/server/continuedev/libs/llm/together.py
new file mode 100644
index 00000000..35b3a424
--- /dev/null
+++ b/server/continuedev/libs/llm/together.py
@@ -0,0 +1,125 @@
+import json
+from typing import Callable
+
+import aiohttp
+from pydantic import Field
+
+from ...core.main import ContinueCustomException
+from ..util.logging import logger
+from .base import LLM
+from .prompts.chat import llama2_template_messages
+from .prompts.edit import simplified_edit_prompt
+
+
+class TogetherLLM(LLM):
+    """
+    The Together API is a cloud platform for running large AI models. You can sign up [here](https://api.together.xyz/signup), copy your API key on the initial welcome screen, and then hit the play button on any model from the [Together Models list](https://docs.together.ai/docs/models-inference). Change `~/.continue/config.py` to look like this:
+
+    ```python title="~/.continue/config.py"
+    from continuedev.core.models import Models
+    from continuedev.libs.llm.together import TogetherLLM
+
+    config = ContinueConfig(
+        ...
+        models=Models(
+            default=TogetherLLM(
+                api_key="<API_KEY>",
+                model="togethercomputer/llama-2-13b-chat"
+            )
+        )
+    )
+    ```
+    """
+
+    api_key: str = Field(..., description="Together API key")
+
+    model: str = "togethercomputer/RedPajama-INCITE-7B-Instruct"
+    base_url: str = Field(
+        "https://api.together.xyz",
+        description="The base URL for your Together API instance",
+    )
+
+    _client_session: aiohttp.ClientSession = None
+
+    template_messages: Callable = llama2_template_messages
+
+    prompt_templates = {
+        "edit": simplified_edit_prompt,
+    }
+
+    async def start(self, **kwargs):
+        await super().start(**kwargs)
+        self._client_session = aiohttp.ClientSession(
+            connector=aiohttp.TCPConnector(verify_ssl=self.verify_ssl),
+            timeout=aiohttp.ClientTimeout(total=self.timeout),
+        )
+
+    async def stop(self):
+        await self._client_session.close()
+
+    async def _stream_complete(self, prompt, options):
+        args = self.collect_args(options)
+
+        async with self._client_session.post(
+            f"{self.base_url}/inference",
+            json={
+                "prompt": prompt,
+                "stream_tokens": True,
+                **args,
+            },
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            proxy=self.proxy,
+        ) as resp:
+            async for line in resp.content.iter_chunks():
+                if line[1]:
+                    json_chunk = line[0].decode("utf-8")
+                    if json_chunk.startswith(": ping - ") or json_chunk.startswith(
+                        "data: [DONE]"
+                    ):
+                        continue
+
+                    chunks = json_chunk.split("\n")
+                    for chunk in chunks:
+                        if chunk.strip() != "":
+                            if chunk.startswith("data: "):
+                                chunk = chunk[6:]
+                            if chunk == "[DONE]":
+                                break
+                            try:
+                                json_chunk = json.loads(chunk)
+                            except Exception as e:
+                                logger.warning(f"Invalid JSON chunk: {chunk}\n\n{e}")
+                                continue
+                            if "choices" in json_chunk:
+                                yield json_chunk["choices"][0]["text"]
+
+    async def _complete(self, prompt: str, options):
+        args = self.collect_args(options)
+
+        async with self._client_session.post(
+            f"{self.base_url}/inference",
+            json={"prompt": prompt, **args},
+            headers={"Authorization": f"Bearer {self.api_key}"},
+            proxy=self.proxy,
+        ) as resp:
+            text = await resp.text()
+            j = json.loads(text)
+            try:
+                if "choices" not in j["output"]:
+                    raise Exception(text)
+                if "output" in j:
+                    return j["output"]["choices"][0]["text"]
+            except Exception as e:
+                j = await resp.json()
+                if "error" in j:
+                    if j["error"].startswith("invalid hexlify value"):
+                        raise ContinueCustomException(
+                            message=f"Invalid Together API key:\n\n{j['error']}",
+                            title="Together API Error",
+                        )
+                    else:
+                        raise ContinueCustomException(
+                            message=j["error"], title="Together API Error"
+                        )
+
+                raise e