feat(openai): add azure support

kharvd · Nov 11, 2024 · edb060b · edb060b
1 parent 5f686f9
commit edb060b
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 2 deletions.
diff --git a/gptcli/assistant.py b/gptcli/assistant.py
@@ -10,6 +10,7 @@
     ModelOverrides,
     Message,
 )
+import gptcli
 from gptcli.providers.google import GoogleCompletionProvider
 from gptcli.providers.llama import LLaMACompletionProvider
 from gptcli.providers.openai import OpenAICompletionProvider
@@ -83,6 +84,8 @@ def get_completion_provider(model: str) -> CompletionProvider:
         return CohereCompletionProvider()
     elif model.startswith("gemini"):
         return GoogleCompletionProvider()
+    elif gptcli.providers.openai.use_azure:
+        return OpenAICompletionProvider()
     else:
         raise ValueError(f"Unknown model: {model}")
 

diff --git a/gptcli/config.py b/gptcli/config.py
@@ -21,6 +21,10 @@ class GptCliConfig:
     api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
     openai_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
     openai_base_url: Optional[str] = os.environ.get("OPENAI_BASE_URL")
+    # When using azure open ai, set your assistant's model to your deployment
+    # name
+    openai_use_azure: bool = False
+    openai_azure_api_version: str = "2024-10-21"
     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")

diff --git a/gptcli/gpt.py b/gptcli/gpt.py
@@ -179,8 +179,15 @@ def main():
     if config.openai_base_url:
         openai.base_url = config.openai_base_url
 
+    if config.openai_use_azure:
+        gptcli.providers.openai.use_azure = config.openai_use_azure
+
+    if config.openai_azure_api_version:
+        openai.api_version = config.openai_azure_api_version
+
     if config.api_key:
         openai.api_key = config.api_key
+
     elif config.openai_api_key:
         openai.api_key = config.openai_api_key
 

diff --git a/gptcli/providers/openai.py b/gptcli/providers/openai.py
@@ -1,7 +1,7 @@
 import re
 from typing import Iterator, List, Optional, cast
 import openai
-from openai import OpenAI
+from openai import AzureOpenAI, OpenAI
 from openai.types.chat import ChatCompletionMessageParam
 
 from gptcli.completion import (
@@ -15,10 +15,16 @@
     UsageEvent,
 )
 
+use_azure: bool = False
 
 class OpenAICompletionProvider(CompletionProvider):
     def __init__(self):
-        self.client = OpenAI(api_key=openai.api_key, base_url=openai.base_url)
+        if use_azure:
+            self.client = AzureOpenAI(api_key=openai.api_key, base_url=openai.base_url, api_version=openai.api_version)
+        else:
+            self.client = OpenAI(api_key=openai.api_key, base_url=openai.base_url)
+
+
 
     def complete(
         self, messages: List[Message], args: dict, stream: bool = False