From 94202717f04cb7721f209934e8669980c38f4501 Mon Sep 17 00:00:00 2001 From: Will Handley Date: Sun, 19 Jan 2025 09:52:53 +0000 Subject: [PATCH 1/6] o1 suggestions --- gptcli/assistant.py | 3 ++ gptcli/config.py | 1 + gptcli/gpt.py | 3 ++ gptcli/providers/grok.py | 107 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 114 insertions(+) create mode 100644 gptcli/providers/grok.py diff --git a/gptcli/assistant.py b/gptcli/assistant.py index 84fb680..ee829b9 100644 --- a/gptcli/assistant.py +++ b/gptcli/assistant.py @@ -15,6 +15,7 @@ from gptcli.providers.anthropic import AnthropicCompletionProvider from gptcli.providers.cohere import CohereCompletionProvider from gptcli.providers.azure_openai import AzureOpenAICompletionProvider +from gptcli.providers.grok import GrokCompletionProvider class AssistantConfig(TypedDict, total=False): @@ -85,6 +86,8 @@ def get_completion_provider(model: str) -> CompletionProvider: return CohereCompletionProvider() elif model.startswith("gemini"): return GoogleCompletionProvider() + elif model.startswith("grok"): + return GrokCompletionProvider() else: raise ValueError(f"Unknown model: {model}") diff --git a/gptcli/config.py b/gptcli/config.py index df5c065..016a724 100644 --- a/gptcli/config.py +++ b/gptcli/config.py @@ -25,6 +25,7 @@ class GptCliConfig: anthropic_api_key: Optional[str] = os.environ.get("ANTHROPIC_API_KEY") google_api_key: Optional[str] = os.environ.get("GOOGLE_API_KEY") cohere_api_key: Optional[str] = os.environ.get("COHERE_API_KEY") + grok_api_key: Optional[str] = os.environ.get("XAI_API_KEY") log_file: Optional[str] = None log_level: str = "INFO" assistants: Dict[str, AssistantConfig] = {} diff --git a/gptcli/gpt.py b/gptcli/gpt.py index 1e55c1f..bc6f4c9 100755 --- a/gptcli/gpt.py +++ b/gptcli/gpt.py @@ -193,6 +193,9 @@ def main(): if config.cohere_api_key: gptcli.providers.cohere.api_key = config.cohere_api_key + if config.grok_api_key: + gptcli.providers.grok.api_key = config.grok_api_key + if config.google_api_key: genai.configure(api_key=config.google_api_key) diff --git a/gptcli/providers/grok.py b/gptcli/providers/grok.py new file mode 100644 index 0000000..c6fdfdb --- /dev/null +++ b/gptcli/providers/grok.py @@ -0,0 +1,107 @@ +import os +import openai +from openai import OpenAI + +from typing import Iterator, List, Optional, cast +from openai.types.chat import ChatCompletionMessageParam +from gptcli.completion import ( + CompletionEvent, + CompletionProvider, + Message, + CompletionError, + BadRequestError, + MessageDeltaEvent, + Pricing, + UsageEvent, +) + + +class GrokCompletionProvider(CompletionProvider): + def __init__(self): + self.api_key = os.environ.get("XAI_API_KEY") or openai.api_key + if not self.api_key: + raise ValueError("XAI_API_KEY environment variable not set and openai.api_key not set") + self.base_url = "https://api.x.ai/v1" + + def complete( + self, messages: List[Message], args: dict, stream: bool = False + ) -> Iterator[CompletionEvent]: + # Save old openai values + old_api_key = openai.api_key + old_base_url = openai.base_url + # Set openai api_key and base_url temporarily + openai.api_key = self.api_key + openai.base_url = self.base_url + + try: + kwargs = {} + if "temperature" in args: + kwargs["temperature"] = args["temperature"] + if "top_p" in args: + kwargs["top_p"] = args["top_p"] + + model = args["model"] + + client = OpenAI( + api_key=self.api_key, + base_url=self.base_url, + ) + + if stream: + response_iter = client.chat.completions.create( + messages=cast(List[ChatCompletionMessageParam], messages), + stream=True, + model=model, + stream_options={"include_usage": True}, + **kwargs, + ) + + for response in response_iter: + if ( + len(response.choices) > 0 + and response.choices[0].finish_reason is None + and response.choices[0].delta.content + ): + yield MessageDeltaEvent(response.choices[0].delta.content) + + if response.usage and (pricing := grok_pricing(args["model"])): + yield UsageEvent.with_pricing( + prompt_tokens=response.usage.prompt_tokens, + completion_tokens=response.usage.completion_tokens, + total_tokens=response.usage.total_tokens, + pricing=pricing, + ) + else: + response = client.chat.completions.create( + messages=messages, + model=model, + stream=False, + **kwargs, + ) + next_choice = response.choices[0] + if next_choice.message.content: + yield MessageDeltaEvent(next_choice.message.content) + if response.usage and (pricing := grok_pricing(args["model"])): + yield UsageEvent.with_pricing( + prompt_tokens=response.usage.prompt_tokens, + completion_tokens=response.usage.completion_tokens, + total_tokens=response.usage.total_tokens, + pricing=pricing, + ) + + except openai.error.InvalidRequestError as e: + raise BadRequestError(str(e)) from e + except openai.error.OpenAIError as e: + raise CompletionError(str(e)) from e + finally: + # Restore old openai values + openai.api_key = old_api_key + openai.base_url = old_base_url + + +def grok_pricing(model: str) -> Optional[Pricing]: + if model.startswith("grok-beta"): + return { + "prompt": 5.00 / 1_000_000, + "response": 15.00 / 1_000_000, + } From 4da5a829aa3f5a81df4403bfdc7761ec45aee95d Mon Sep 17 00:00:00 2001 From: Will Handley Date: Sun, 19 Jan 2025 09:56:36 +0000 Subject: [PATCH 2/6] Changed name to xai --- gptcli/assistant.py | 4 ++-- gptcli/config.py | 2 +- gptcli/gpt.py | 4 ++-- gptcli/providers/{grok.py => xai.py} | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) rename gptcli/providers/{grok.py => xai.py} (98%) diff --git a/gptcli/assistant.py b/gptcli/assistant.py index ee829b9..e08f3f7 100644 --- a/gptcli/assistant.py +++ b/gptcli/assistant.py @@ -15,7 +15,7 @@ from gptcli.providers.anthropic import AnthropicCompletionProvider from gptcli.providers.cohere import CohereCompletionProvider from gptcli.providers.azure_openai import AzureOpenAICompletionProvider -from gptcli.providers.grok import GrokCompletionProvider +from gptcli.providers.xai import XAICompletionProvider class AssistantConfig(TypedDict, total=False): @@ -87,7 +87,7 @@ def get_completion_provider(model: str) -> CompletionProvider: elif model.startswith("gemini"): return GoogleCompletionProvider() elif model.startswith("grok"): - return GrokCompletionProvider() + return XAICompletionProvider() else: raise ValueError(f"Unknown model: {model}") diff --git a/gptcli/config.py b/gptcli/config.py index 016a724..0d5ce15 100644 --- a/gptcli/config.py +++ b/gptcli/config.py @@ -25,7 +25,7 @@ class GptCliConfig: anthropic_api_key: Optional[str] = os.environ.get("ANTHROPIC_API_KEY") google_api_key: Optional[str] = os.environ.get("GOOGLE_API_KEY") cohere_api_key: Optional[str] = os.environ.get("COHERE_API_KEY") - grok_api_key: Optional[str] = os.environ.get("XAI_API_KEY") + xai_api_key: Optional[str] = os.environ.get("XAI_API_KEY") log_file: Optional[str] = None log_level: str = "INFO" assistants: Dict[str, AssistantConfig] = {} diff --git a/gptcli/gpt.py b/gptcli/gpt.py index bc6f4c9..ec8faba 100755 --- a/gptcli/gpt.py +++ b/gptcli/gpt.py @@ -193,8 +193,8 @@ def main(): if config.cohere_api_key: gptcli.providers.cohere.api_key = config.cohere_api_key - if config.grok_api_key: - gptcli.providers.grok.api_key = config.grok_api_key + if config.xai_api_key: + gptcli.providers.xai.api_key = config.xai_api_key if config.google_api_key: genai.configure(api_key=config.google_api_key) diff --git a/gptcli/providers/grok.py b/gptcli/providers/xai.py similarity index 98% rename from gptcli/providers/grok.py rename to gptcli/providers/xai.py index c6fdfdb..a7136f0 100644 --- a/gptcli/providers/grok.py +++ b/gptcli/providers/xai.py @@ -16,7 +16,7 @@ ) -class GrokCompletionProvider(CompletionProvider): +class XAICompletionProvider(CompletionProvider): def __init__(self): self.api_key = os.environ.get("XAI_API_KEY") or openai.api_key if not self.api_key: From 4b9b4ad1f0b19b66af2da380d4824924c4d20514 Mon Sep 17 00:00:00 2001 From: Will Handley Date: Sun, 19 Jan 2025 10:23:43 +0000 Subject: [PATCH 3/6] Trimmed down and re-used classes --- gptcli/providers/openai.py | 51 +++++++-------- gptcli/providers/xai.py | 125 +++++++------------------------------ 2 files changed, 50 insertions(+), 126 deletions(-) diff --git a/gptcli/providers/openai.py b/gptcli/providers/openai.py index 14a8035..c716dbd 100644 --- a/gptcli/providers/openai.py +++ b/gptcli/providers/openai.py @@ -54,7 +54,7 @@ def complete( ): yield MessageDeltaEvent(response.choices[0].delta.content) - if response.usage and (pricing := gpt_pricing(args["model"])): + if response.usage and (pricing := self.pricing(args["model"])): yield UsageEvent.with_pricing( prompt_tokens=response.usage.prompt_tokens, completion_tokens=response.usage.completion_tokens, @@ -71,7 +71,7 @@ def complete( next_choice = response.choices[0] if next_choice.message.content: yield MessageDeltaEvent(next_choice.message.content) - if response.usage and (pricing := gpt_pricing(args["model"])): + if response.usage and (pricing := self.pricing(args["model"])): yield UsageEvent.with_pricing( prompt_tokens=response.usage.prompt_tokens, completion_tokens=response.usage.completion_tokens, @@ -84,6 +84,30 @@ def complete( except openai.APIError as e: raise CompletionError(e.message) from e + def pricing(self, model: str) -> Optional[Pricing]: + if model.startswith("gpt-3.5-turbo-16k"): + return GPT_3_5_TURBO_16K_PRICE_PER_TOKEN + elif model.startswith("gpt-3.5-turbo"): + return GPT_3_5_TURBO_PRICE_PER_TOKEN + elif model.startswith("gpt-4-32k"): + return GPT_4_32K_PRICE_PER_TOKEN + elif model.startswith("gpt-4o-mini"): + return GPT_4_O_MINI_PRICE_PER_TOKEN + elif model.startswith("gpt-4o-2024-05-13") or model.startswith("chatgpt-4o-latest"): + return GPT_4_O_2024_05_13_PRICE_PER_TOKEN + elif model.startswith("gpt-4o"): + return GPT_4_O_2024_08_06_PRICE_PER_TOKEN + elif model.startswith("gpt-4-turbo") or re.match(r"gpt-4-\d\d\d\d-preview", model): + return GPT_4_TURBO_PRICE_PER_TOKEN + elif model.startswith("gpt-4"): + return GPT_4_PRICE_PER_TOKEN + elif model.startswith("o1-preview"): + return O_1_PREVIEW_PRICE_PER_TOKEN + elif model.startswith("o1-mini"): + return O_1_MINI_PRICE_PER_TOKEN + else: + return None + GPT_3_5_TURBO_PRICE_PER_TOKEN: Pricing = { "prompt": 0.50 / 1_000_000, @@ -135,26 +159,3 @@ def complete( "response": 12.0 / 1_000_000, } -def gpt_pricing(model: str) -> Optional[Pricing]: - if model.startswith("gpt-3.5-turbo-16k"): - return GPT_3_5_TURBO_16K_PRICE_PER_TOKEN - elif model.startswith("gpt-3.5-turbo"): - return GPT_3_5_TURBO_PRICE_PER_TOKEN - elif model.startswith("gpt-4-32k"): - return GPT_4_32K_PRICE_PER_TOKEN - elif model.startswith("gpt-4o-mini"): - return GPT_4_O_MINI_PRICE_PER_TOKEN - elif model.startswith("gpt-4o-2024-05-13") or model.startswith("chatgpt-4o-latest"): - return GPT_4_O_2024_05_13_PRICE_PER_TOKEN - elif model.startswith("gpt-4o"): - return GPT_4_O_2024_08_06_PRICE_PER_TOKEN - elif model.startswith("gpt-4-turbo") or re.match(r"gpt-4-\d\d\d\d-preview", model): - return GPT_4_TURBO_PRICE_PER_TOKEN - elif model.startswith("gpt-4"): - return GPT_4_PRICE_PER_TOKEN - elif model.startswith("o1-preview"): - return O_1_PREVIEW_PRICE_PER_TOKEN - elif model.startswith("o1-mini"): - return O_1_MINI_PRICE_PER_TOKEN - else: - return None diff --git a/gptcli/providers/xai.py b/gptcli/providers/xai.py index a7136f0..98d3a52 100644 --- a/gptcli/providers/xai.py +++ b/gptcli/providers/xai.py @@ -1,107 +1,30 @@ import os -import openai from openai import OpenAI +from typing import Optional +from gptcli.providers.openai import OpenAICompletionProvider +from gptcli.completion import Pricing -from typing import Iterator, List, Optional, cast -from openai.types.chat import ChatCompletionMessageParam -from gptcli.completion import ( - CompletionEvent, - CompletionProvider, - Message, - CompletionError, - BadRequestError, - MessageDeltaEvent, - Pricing, - UsageEvent, -) +api_key = os.environ.get("XAI_API_KEY") - -class XAICompletionProvider(CompletionProvider): +class XAICompletionProvider(OpenAICompletionProvider): def __init__(self): - self.api_key = os.environ.get("XAI_API_KEY") or openai.api_key - if not self.api_key: - raise ValueError("XAI_API_KEY environment variable not set and openai.api_key not set") - self.base_url = "https://api.x.ai/v1" - - def complete( - self, messages: List[Message], args: dict, stream: bool = False - ) -> Iterator[CompletionEvent]: - # Save old openai values - old_api_key = openai.api_key - old_base_url = openai.base_url - # Set openai api_key and base_url temporarily - openai.api_key = self.api_key - openai.base_url = self.base_url - - try: - kwargs = {} - if "temperature" in args: - kwargs["temperature"] = args["temperature"] - if "top_p" in args: - kwargs["top_p"] = args["top_p"] - - model = args["model"] - - client = OpenAI( - api_key=self.api_key, - base_url=self.base_url, - ) - - if stream: - response_iter = client.chat.completions.create( - messages=cast(List[ChatCompletionMessageParam], messages), - stream=True, - model=model, - stream_options={"include_usage": True}, - **kwargs, - ) - - for response in response_iter: - if ( - len(response.choices) > 0 - and response.choices[0].finish_reason is None - and response.choices[0].delta.content - ): - yield MessageDeltaEvent(response.choices[0].delta.content) - - if response.usage and (pricing := grok_pricing(args["model"])): - yield UsageEvent.with_pricing( - prompt_tokens=response.usage.prompt_tokens, - completion_tokens=response.usage.completion_tokens, - total_tokens=response.usage.total_tokens, - pricing=pricing, - ) - else: - response = client.chat.completions.create( - messages=messages, - model=model, - stream=False, - **kwargs, - ) - next_choice = response.choices[0] - if next_choice.message.content: - yield MessageDeltaEvent(next_choice.message.content) - if response.usage and (pricing := grok_pricing(args["model"])): - yield UsageEvent.with_pricing( - prompt_tokens=response.usage.prompt_tokens, - completion_tokens=response.usage.completion_tokens, - total_tokens=response.usage.total_tokens, - pricing=pricing, - ) - - except openai.error.InvalidRequestError as e: - raise BadRequestError(str(e)) from e - except openai.error.OpenAIError as e: - raise CompletionError(str(e)) from e - finally: - # Restore old openai values - openai.api_key = old_api_key - openai.base_url = old_base_url - + self.client = OpenAI(api_key=api_key, base_url="https://api.x.ai/v1") + + def pricing(self, model: str) -> Optional[Pricing]: + if model.startswith("grok-beta"): + return GROK_BETA_PRICE_PER_TOKEN + elif model.startswith("grok-2"): + return GROK_2_PRICE_PER_TOKEN + else: + return None + +GROK_2_PRICE_PER_TOKEN: Pricing = { + "prompt": 2.00 / 1_000_000, + "response": 10.00 / 1_000_000 +} + +GROK_BETA_PRICE_PER_TOKEN: Pricing = { + "prompt": 5.00 / 1_000_000, + "response": 15.00 / 1_000_000 +} -def grok_pricing(model: str) -> Optional[Pricing]: - if model.startswith("grok-beta"): - return { - "prompt": 5.00 / 1_000_000, - "response": 15.00 / 1_000_000, - } From 4dde1ddc2318d3707b2184c4bbe282692db0e922 Mon Sep 17 00:00:00 2001 From: Will Handley Date: Sun, 19 Jan 2025 10:31:52 +0000 Subject: [PATCH 4/6] re-instated gpt_pricing function for backward compatibility --- gptcli/providers/openai.py | 47 +++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/gptcli/providers/openai.py b/gptcli/providers/openai.py index c716dbd..10717e5 100644 --- a/gptcli/providers/openai.py +++ b/gptcli/providers/openai.py @@ -85,29 +85,7 @@ def complete( raise CompletionError(e.message) from e def pricing(self, model: str) -> Optional[Pricing]: - if model.startswith("gpt-3.5-turbo-16k"): - return GPT_3_5_TURBO_16K_PRICE_PER_TOKEN - elif model.startswith("gpt-3.5-turbo"): - return GPT_3_5_TURBO_PRICE_PER_TOKEN - elif model.startswith("gpt-4-32k"): - return GPT_4_32K_PRICE_PER_TOKEN - elif model.startswith("gpt-4o-mini"): - return GPT_4_O_MINI_PRICE_PER_TOKEN - elif model.startswith("gpt-4o-2024-05-13") or model.startswith("chatgpt-4o-latest"): - return GPT_4_O_2024_05_13_PRICE_PER_TOKEN - elif model.startswith("gpt-4o"): - return GPT_4_O_2024_08_06_PRICE_PER_TOKEN - elif model.startswith("gpt-4-turbo") or re.match(r"gpt-4-\d\d\d\d-preview", model): - return GPT_4_TURBO_PRICE_PER_TOKEN - elif model.startswith("gpt-4"): - return GPT_4_PRICE_PER_TOKEN - elif model.startswith("o1-preview"): - return O_1_PREVIEW_PRICE_PER_TOKEN - elif model.startswith("o1-mini"): - return O_1_MINI_PRICE_PER_TOKEN - else: - return None - + return gpt_pricing(model) GPT_3_5_TURBO_PRICE_PER_TOKEN: Pricing = { "prompt": 0.50 / 1_000_000, @@ -159,3 +137,26 @@ def pricing(self, model: str) -> Optional[Pricing]: "response": 12.0 / 1_000_000, } +def gpt_pricing(model: str) -> Optional[Pricing]: + if model.startswith("gpt-3.5-turbo-16k"): + return GPT_3_5_TURBO_16K_PRICE_PER_TOKEN + elif model.startswith("gpt-3.5-turbo"): + return GPT_3_5_TURBO_PRICE_PER_TOKEN + elif model.startswith("gpt-4-32k"): + return GPT_4_32K_PRICE_PER_TOKEN + elif model.startswith("gpt-4o-mini"): + return GPT_4_O_MINI_PRICE_PER_TOKEN + elif model.startswith("gpt-4o-2024-05-13") or model.startswith("chatgpt-4o-latest"): + return GPT_4_O_2024_05_13_PRICE_PER_TOKEN + elif model.startswith("gpt-4o"): + return GPT_4_O_2024_08_06_PRICE_PER_TOKEN + elif model.startswith("gpt-4-turbo") or re.match(r"gpt-4-\d\d\d\d-preview", model): + return GPT_4_TURBO_PRICE_PER_TOKEN + elif model.startswith("gpt-4"): + return GPT_4_PRICE_PER_TOKEN + elif model.startswith("o1-preview"): + return O_1_PREVIEW_PRICE_PER_TOKEN + elif model.startswith("o1-mini"): + return O_1_MINI_PRICE_PER_TOKEN + else: + return None From e63b50d0fc490a47ac08021cee91e4f4a4865176 Mon Sep 17 00:00:00 2001 From: Will Handley Date: Sun, 19 Jan 2025 10:32:04 +0000 Subject: [PATCH 5/6] black and isort --- gptcli/providers/xai.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/gptcli/providers/xai.py b/gptcli/providers/xai.py index 98d3a52..b47f0e6 100644 --- a/gptcli/providers/xai.py +++ b/gptcli/providers/xai.py @@ -1,11 +1,14 @@ import os -from openai import OpenAI from typing import Optional -from gptcli.providers.openai import OpenAICompletionProvider + +from openai import OpenAI + from gptcli.completion import Pricing +from gptcli.providers.openai import OpenAICompletionProvider api_key = os.environ.get("XAI_API_KEY") + class XAICompletionProvider(OpenAICompletionProvider): def __init__(self): self.client = OpenAI(api_key=api_key, base_url="https://api.x.ai/v1") @@ -18,13 +21,13 @@ def pricing(self, model: str) -> Optional[Pricing]: else: return None + GROK_2_PRICE_PER_TOKEN: Pricing = { "prompt": 2.00 / 1_000_000, - "response": 10.00 / 1_000_000 + "response": 10.00 / 1_000_000, } GROK_BETA_PRICE_PER_TOKEN: Pricing = { "prompt": 5.00 / 1_000_000, - "response": 15.00 / 1_000_000 + "response": 15.00 / 1_000_000, } - From 7279cf107c4e17c14f6c3d4a24c61c98a249d527 Mon Sep 17 00:00:00 2001 From: Will Handley Date: Sun, 19 Jan 2025 10:46:36 +0000 Subject: [PATCH 6/6] Bumped patch version number --- gptcli/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gptcli/__init__.py b/gptcli/__init__.py index 493f741..260c070 100644 --- a/gptcli/__init__.py +++ b/gptcli/__init__.py @@ -1 +1 @@ -__version__ = "0.3.0" +__version__ = "0.3.1"