Add together.xyz endpoint

gptcli/assistant.py

@@ -9,13 +9,19 @@
 from gptcli.llama import LLaMACompletionProvider
 from gptcli.openai import OpenAICompletionProvider
 from gptcli.anthropic import AnthropicCompletionProvider
+from gptcli.together import TogetherCompletionProvider
 
 
 class AssistantConfig(TypedDict, total=False):
     messages: List[Message]
     model: str
     temperature: float
     top_p: float
+    system_prefix: str
+    system_suffix: str
+    user_prefix: str
+    user_suffix: str
+    stop_tokens: List[str]
 
 
 CONFIG_DEFAULTS = {
@@ -55,7 +61,9 @@ class AssistantConfig(TypedDict, total=False):
 }
 
 
-def get_completion_provider(model: str) -> CompletionProvider:
+def get_completion_provider(
+    model: str, assistant_config: AssistantConfig
+) -> CompletionProvider:
     if model.startswith("gpt"):
         return OpenAICompletionProvider()
     elif model.startswith("claude"):
@@ -64,6 +72,16 @@ def get_completion_provider(model: str) -> CompletionProvider:
         return LLaMACompletionProvider()
     elif model.startswith("chat-bison"):
         return GoogleCompletionProvider()
+    elif model.startswith("together"):
+        return TogetherCompletionProvider(
+            {
+                "system_prefix": assistant_config.get("system_prefix", ""),
+                "system_suffix": assistant_config.get("system_suffix", ""),
+                "user_prefix": assistant_config.get("user_prefix", ""),
+                "user_suffix": assistant_config.get("user_suffix", ""),
+                "stop_tokens": assistant_config.get("stop_tokens", []),
+            }
+        )
     else:
         raise ValueError(f"Unknown model: {model}")
 
@@ -103,7 +121,7 @@ def complete_chat(
         self, messages, override_params: ModelOverrides = {}, stream: bool = True
     ) -> Iterator[str]:
         model = self._param("model", override_params)
-        completion_provider = get_completion_provider(model)
+        completion_provider = get_completion_provider(model, self.config)
         return completion_provider.complete(
             messages,
             {

gptcli/config.py

@@ -22,6 +22,7 @@ class GptCliConfig:
     openai_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
     anthropic_api_key: Optional[str] = os.environ.get("ANTHROPIC_API_KEY")
     google_api_key: Optional[str] = os.environ.get("GOOGLE_API_KEY")
+    together_api_key: Optional[str] = os.environ.get("TOGETHER_API_KEY")
     log_file: Optional[str] = None
     log_level: str = "INFO"
     assistants: Dict[str, AssistantConfig] = {}

gptcli/gpt.py

@@ -15,6 +15,7 @@
 import datetime
 import google.generativeai as genai
 import gptcli.anthropic
+import gptcli.together
 from gptcli.assistant import (
     Assistant,
     DEFAULT_ASSISTANTS,
@@ -178,6 +179,9 @@ def main():
         )
         sys.exit(1)
 
+    if config.together_api_key:
+        gptcli.together.api_key = config.together_api_key
+
     if config.anthropic_api_key:
         gptcli.anthropic.api_key = config.anthropic_api_key
 

gptcli/together.py

@@ -0,0 +1,84 @@
+import json
+import logging
+import os
+import requests
+import sseclient
+
+from typing import Iterator, List, TypedDict, cast
+
+from gptcli.completion import CompletionProvider, Message
+
+api_key = os.environ.get("TOGETHER_API_KEY")
+url = "https://api.together.xyz/inference"
+
+
+class PromptConfig(TypedDict):
+    system_prefix: str
+    system_suffix: str
+    user_prefix: str
+    user_suffix: str
+    stop_tokens: List[str]
+
+
+def build_prompt(messages: List[Message], prompt_config: PromptConfig) -> str:
+    prompt = ""
+    for message in messages:
+        if message["role"] == "system":
+            prompt += prompt_config["system_prefix"]
+            prompt += message["content"]
+            prompt += prompt_config["system_suffix"]
+        elif message["role"] == "user":
+            prompt += prompt_config["user_prefix"]
+            prompt += message["content"]
+            prompt += prompt_config["user_suffix"]
+        else:
+            prompt += message["content"]
+    return prompt
+
+
+class TogetherCompletionProvider(CompletionProvider):
+    def __init__(self, prompt_config: PromptConfig):
+        self.prompt_config = prompt_config
+
+    def complete(
+        self, messages: List[Message], args: dict, stream: bool = False
+    ) -> Iterator[str]:
+        kwargs = {}
+        if "temperature" in args:
+            kwargs["temperature"] = args["temperature"]
+        if "top_p" in args:
+            kwargs["top_p"] = args["top_p"]
+
+        assert stream, "Together only supports streaming completions"
+
+        model = args["model"].split("/", 1)[1]
+        prompt = build_prompt(messages, self.prompt_config)
+
+        logging.info(f"Prompt: {prompt}")
+
+        payload = {
+            "model": model,
+            "prompt": prompt,
+            "max_tokens": 2048,
+            "stream_tokens": True,
+            "stop": self.prompt_config["stop_tokens"],
+            **kwargs,
+        }
+
+        headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+
+        response = requests.post(url, json=payload, headers=headers, stream=True)
+        response.raise_for_status()
+
+        client = sseclient.SSEClient(response)
+        for event in client.events():
+            if event.data == "[DONE]":
+                break
+
+            partial_result = json.loads(event.data)
+            token = partial_result["choices"][0]["text"]
+            yield token

pyproject.toml

@@ -26,6 +26,8 @@ dependencies = [
     "pytest==7.3.1",
     "PyYAML==6.0",
     "rich==13.7.0",
+    "requests==2.31.0",
+    "sseclient-py==1.8.0",
     "tiktoken==0.5.2",
     "tokenizers==0.15.0",
     "typing_extensions==4.5.0",