adding support for custom openai client for Azure use

Author: DARREN OBERST

examples/Models/using-azure-openai.py

@@ -0,0 +1,51 @@
+
+""" This example shows how to use OpenAIConfigs to create a configured OpenAI client, most often used for
+Azure OpenAI access."""
+
+import os
+
+from llmware.models import ModelCatalog
+from llmware.configs import OpenAIConfig
+from openai import AzureOpenAI, OpenAI
+
+
+#   to start - OpenAI client is created in OpenAI Generative and Embedding models classes at the time of inference
+#   the client will be created as a standard OpenAI client with the api_keys passed
+
+my_azure_client = OpenAIConfig().get_azure_client()
+print("my azure client to start: ", my_azure_client)
+
+#   to configure an AzureOpenAI client, two steps:
+#   first, create the client with openai >= 1.0 python SDK, (see above) e.g.:
+
+client = AzureOpenAI(
+  azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"),
+  api_key=os.getenv("AZURE_OPENAI_API_KEY"),
+  api_version="2024-02-01"
+)
+
+#   second, set the azure client in OpenAIConfigs as below:
+OpenAIConfig().set_azure_client(client)
+print("my azure client - set: ", OpenAIConfig().get_azure_client())
+
+#   now, run the inference like any other in llmware
+
+#   OpenAI Generative call
+model = ModelCatalog().load_model("gpt-4")
+
+#   the model will check the value of get_azure_client() in the configs -> if set, then will use
+response = model.inference("What is the future of AI")
+print("response: ", response)
+
+#   OpenAI Embedding call
+model = ModelCatalog().load_model("text-embedding-3-small")
+embedding = model.embedding(["This is a sample sentence for an embedding test."])
+print("embedding: ", embedding)
+
+#   reset so you can use the standard OpenAI client
+OpenAIConfig().set_azure_client(None)
+
+model = ModelCatalog().load_model("text-embedding-3-small", api_key="your openai api key")
+embedding = model.embedding(["This is a sample sentence for an embedding test."])
+print("embedding: ", embedding)
+

llmware/configs.py

@@ -896,3 +896,58 @@ def get_auth_credentials(cls):
     @classmethod
     def get_auth_token_transport_header(cls):
         return cls._conf["auth_token_transport_header"]
+
+
+class OpenAIConfig:
+
+    """Configuration object for OpenAI - primarily for configuring Azure OpenAI credentials.
+
+     Primary use is to setup an AzureOpenAI client that will be used in place of the standard
+     OpenAI client.
+
+     Within LLMWare, the OpenAI model classes will check this config before creating a new OpenAI client.
+
+     If there is a client already established in _conf["openai_client"], that client will be used in the
+     inference/embedding process.
+
+     For example:
+
+        # create your AzureOpenAI client
+
+        from openai import AzureOpenAI
+
+        client = AzureOpenAI(
+                            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+                            api_key=os.getenv("AZURE_OPENAI_API_KEY"),
+                            api_version="2024-02-01")
+
+        # add that client to the OpenAIConfig:
+
+        OpenAIConfig().set_azure_client(client)
+
+     """
+
+    _conf = {"openai_client": None,
+             "api_key": None,               # placeholder / not used currently
+             "api_version": None,           # placeholder / not used currently
+             "use_azure_endpoint": False}   # placeholder / not used currently
+
+    @classmethod
+    def get_config(cls, name):
+        if name in cls._conf:
+            return cls._conf[name]
+        raise ConfigKeyException(name)
+
+    @classmethod
+    def set_config(cls, name, value):
+        cls._conf[name] = value
+
+    @classmethod
+    def set_azure_client(cls, azure_client):
+        cls._conf["openai_client"] = azure_client
+
+    @classmethod
+    def get_azure_client(cls):
+        return cls._conf["openai_client"]
+
+

llmware/models.py

@@ -2564,6 +2564,9 @@ def inference(self, prompt, add_context=None, add_prompt_engineering=None, infer
             if "max_tokens" in inference_dict:
                 self.target_requested_output_tokens = inference_dict["max_tokens"]
 
+            if "openai_client" in inference_dict:
+                self.openai_client = inference_dict["openai_client"]
+
         # api_key
         if api_key:
             self.api_key = api_key
@@ -2583,6 +2586,8 @@ def inference(self, prompt, add_context=None, add_prompt_engineering=None, infer
         except ImportError:
             raise DependencyNotInstalledException("openai >= 1.0")
 
+        from llmware.configs import OpenAIConfig
+
         usage = {}
         time_start = time.time()
 
@@ -2595,7 +2600,17 @@ def inference(self, prompt, add_context=None, add_prompt_engineering=None, infer
 
                 # updated OpenAI client to >v1.0 API - create client, and returns pydantic objects
 
-                client = OpenAI(api_key=self.api_key)
+                azure_client = OpenAIConfig().get_azure_client()
+
+                if not azure_client:
+                    client = OpenAI(api_key=self.api_key)
+
+                else:
+
+                    logging.info("update: applying custom OpenAI client from OpenAIConfig")
+
+                    client = azure_client
+
                 response = client.chat.completions.create(model=self.model_name,messages=messages,
                                                           max_tokens=self.target_requested_output_tokens)
 
@@ -2616,7 +2631,17 @@ def inference(self, prompt, add_context=None, add_prompt_engineering=None, infer
 
                 text_prompt = prompt_final + self.separator
 
-                client = OpenAI(api_key=self.api_key)
+                azure_client = OpenAIConfig().get_azure_client()
+
+                if not azure_client:
+                    client = OpenAI(api_key=self.api_key)
+
+                else:
+
+                    logging.info("update: applying custom OpenAI client from OpenAIConfig")
+
+                    client = azure_client
+
                 response = client.completions.create(model=self.model_name, prompt=text_prompt,
                                                      temperature=self.temperature,
                                                      max_tokens=self.target_requested_output_tokens)
@@ -3697,6 +3722,8 @@ def embedding(self, text_sample, api_key=None):
         except ImportError:
             raise DependencyNotInstalledException("openai >= 1.0")
 
+        from llmware.configs import OpenAIConfig
+
         # insert safety check here
         safe_samples = []
         safety_buffer = 200
@@ -3732,7 +3759,18 @@ def embedding(self, text_sample, api_key=None):
         # end - safety check
 
         # update to open >v1.0 api - create client and output is pydantic objects
-        client = OpenAI(api_key=self.api_key)
+
+        azure_client = OpenAIConfig().get_azure_client()
+
+        if not azure_client:
+            client = OpenAI(api_key=self.api_key)
+
+        else:
+
+            logging.info("update: applying custom OpenAI client from OpenAIConfig")
+
+            client = azure_client
+
         response = client.embeddings.create(model=self.model_name, input=text_prompt)
 
         # logging.info("update: response: %s ", response)