From 2efbac61a73c2dfef54e0ef4dfee19b5509bce4e Mon Sep 17 00:00:00 2001
From: Valery Kharitonov <val@kharvd.com>
Date: Tue, 14 May 2024 23:31:57 -0400
Subject: [PATCH] Allow non-OpenAI model names with oai-compat: prefix

---
 README.md           | 19 ++++++++++++++++++-
 gptcli/assistant.py |  6 +++++-
 gptcli/openai.py    |  8 ++++++--
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 69d68f6..6950275 100644
--- a/README.md
+++ b/README.md
@@ -21,16 +21,19 @@ Command-line interface for ChatGPT Claude and Bard.
 ## Installation
 
 This install assumes a Linux/OSX machine with Python and pip available.
+
 ```bash
 pip install gpt-command-line
 ```
 
 Install latest version from source:
+
 ```bash
 pip install git+https://github.com/kharvd/gpt-cli.git
 ```
 
 Or install by cloning the repository manually:
+
 ```bash
 git clone https://github.com/kharvd/gpt-cli.git
 cd gpt-cli
@@ -173,13 +176,23 @@ Ahoy, matey! What be bringing ye to these here waters? Be it treasure or adventu
 ```
 
 ### Customize OpenAI API URL
-If you are using other models compatible with the OpenAI Python SDK, you can configure them by modifying the `openai_base_url` setting in the config file or using the `OPENAI_BASE_URL` environment variable .  
+
+If you are using other models compatible with the OpenAI Python SDK, you can configure them by modifying the `openai_base_url` setting in the config file or using the `OPENAI_BASE_URL` environment variable .
 
 Example:
+
 ```
 openai_base_url: https://your-custom-api-url.com/v1
 ```
 
+Use `oai-compat:` prefix for the model name to pass non-GPT model names to the API. For example, to chat with Llama3-70b on [https://together.ai](Together), use the following command:
+
+```bash
+OPENAI_API_KEY=$TOGETHER_API_KEY OPENAI_BASE_URL=https://api.together.xyz/v1 gpt general --model oai-compat:meta-llama/Llama-3-70b-chat-hf
+```
+
+The prefix is stripped before sending the request to the API.
+
 ## Other chat bots
 
 ### Anthropic Claude
@@ -203,17 +216,21 @@ gpt --model claude-v1
 ```
 
 ### Google Bard (PaLM 2)
+
 Similar to Claude, set the Google API key
 
 ```bash
 export GOOGLE_API_KEY=<your_key_here>
 ```
+
 or a config line:
+
 ```yaml
 google_api_key: <your_key_here>
 ```
 
 Run `gpt` with the correct model:
+
 ```bash
 gpt --model chat-bison-001
 ```
diff --git a/gptcli/assistant.py b/gptcli/assistant.py
index dc4df20..61ae9ba 100644
--- a/gptcli/assistant.py
+++ b/gptcli/assistant.py
@@ -66,7 +66,11 @@ class AssistantConfig(TypedDict, total=False):
 
 
 def get_completion_provider(model: str) -> CompletionProvider:
-    if model.startswith("gpt") or model.startswith("ft:gpt"):
+    if (
+        model.startswith("gpt")
+        or model.startswith("ft:gpt")
+        or model.startswith("oai-compat:")
+    ):
         return OpenAICompletionProvider()
     elif model.startswith("claude"):
         return AnthropicCompletionProvider()
diff --git a/gptcli/openai.py b/gptcli/openai.py
index d1b8639..6e64673 100644
--- a/gptcli/openai.py
+++ b/gptcli/openai.py
@@ -29,12 +29,16 @@ def complete(
         if "top_p" in args:
             kwargs["top_p"] = args["top_p"]
 
+        model = args["model"]
+        if model.startswith("oai-compat:"):
+            model = model[len("oai-compat:") :]
+
         try:
             if stream:
                 response_iter = self.client.chat.completions.create(
                     messages=cast(List[ChatCompletionMessageParam], messages),
                     stream=True,
-                    model=args["model"],
+                    model=model,
                     stream_options={"include_usage": True},
                     **kwargs,
                 )
@@ -57,7 +61,7 @@ def complete(
             else:
                 response = self.client.chat.completions.create(
                     messages=cast(List[ChatCompletionMessageParam], messages),
-                    model=args["model"],
+                    model=model,
                     stream=False,
                     **kwargs,
                 )