Update Azure OpenAI use

README.md

@@ -106,6 +106,23 @@ pandas_gpt.completer = pandas_gpt.LiteLLM('huggingface/meta-llama/Meta-Llama-3.1
 pandas_gpt.completer = pandas_gpt.OpenRouter('anthropic/claude-3.5-sonnet')
 ```
 
+### Azure OpenAI
+
+Get your settings from [Azure AI Foundry](https://oai.azure.com/) or your system administrator.
+
+```python
+import openai
+import pandas_gpt
+
+openai.api_key = '<API Key>'
+
+pandas_gpt.completer = pandas_gpt.AzureOpenAI(
+  deployment=<Deployment ID>,
+  azure_endpoint=<Azure OpenAI Endpoint>,
+  api_version=<Api version>,
+)
+```
+
 ### Anything
 
 ```python
@@ -118,24 +135,6 @@ def my_custom_completer(prompt: str) -> str:
 pandas_gpt.completer = my_custom_completer
 ```
 
-If you want to use a fully customized API host such as [Azure OpenAI Service](https://azure.microsoft.com/en-us/products/cognitive-services/openai-service),
-you can globally configure the `openai` and `pandas-gpt` packages:
-
-```python
-import openai
-openai.api_type = 'azure'
-openai.api_base = '<Endpoint>'
-openai.api_version = '<Version>'
-openai.api_key = '<API Key>'
-
-import pandas_gpt
-pandas_gpt.completer = pandas_gpt.OpenAI(
-  model='gpt-3.5-turbo',
-  engine='<Engine>',
-  deployment_id='<Deployment ID>',
-)
-```
-
 ## Alternatives
 
 - [GitHub Copilot](https://github.com/features/copilot): General-purpose code completion (paid subscription)

src/pandas_gpt/__init__.py

@@ -2,7 +2,7 @@
 from typing import Any, Callable
 import pandas as pd
 
-from .completers import LiteLLM, OpenAI, OpenRouter
+from .completers import AzureOpenAI, LiteLLM, OpenAI, OpenRouter
 
 __all__ = [
     # Global config
@@ -20,6 +20,7 @@
     "LiteLLM",
     "OpenAI",
     "OpenRouter",
+    "AzureOpenAI",
 ]
 
 # Override with `pandas_gpt.verbose = True`

src/pandas_gpt/completers/__init__.py

@@ -1,3 +1,4 @@
+from .azureopenai import AzureOpenAI
 from .litellm import LiteLLM
 from .openai import OpenAI
 from .openrouter import OpenRouter

src/pandas_gpt/completers/azureopenai.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass
+import os
+from typing import Any
+
+__all__ = ["AzureOpenAI"]
+
+default_system_prompt = "Write the function in a Python code block with all necessary imports and no example usage."
+
+
+@dataclass
+class AzureOpenAI:
+    completion_config: dict[str, Any]
+    client_config: dict[str, Any]
+    system_prompt: str
+    _cache: dict[str, str]
+    _client: Any
+
+    def __init__(self, deployment: str, **client_config):
+        self.completion_config = {"model": deployment}
+        self.client_config = client_config
+        self.system_prompt = default_system_prompt
+        self._cache = {}
+        self._client = None
+
+    def __call__(self, prompt: str) -> str:
+        completion = self._cache.get(prompt) or self.run_completion_function(
+            **self.completion_config,
+            messages=[
+                dict(role="system", content=self.system_prompt),
+                dict(role="user", content=prompt),
+            ],
+        )
+        self._cache[prompt] = completion
+        return completion.choices[0].message.content
+
+    def run_completion_function(self, **kw):
+        if self._client is None:
+            try:
+                import openai
+            except ImportError:
+                raise Exception(
+                    "The package `openai` could not be found. You can fix this error by running `pip install pandas-gpt[openai]` or passing a custom `completer` argument."
+                )
+            client_config = dict(self.client_config)
+            api_key = os.environ.get("AZURE_OPENAI_API_KEY", openai.api_key)
+            if api_key is not None and "api_key" not in client_config:
+                client_config["api_key"] = api_key
+            self._client = openai.AzureOpenAI(**client_config)
+        return self._client.chat.completions.create(**kw)