Change OpenAI to use Identity Provider (#122)

Author: BenConstable9

deploy_ai_search/src/deploy_ai_search/environment.py

@@ -109,12 +109,11 @@ def ai_search_credential(self) -> DefaultAzureCredential | AzureKeyCredential:
         Returns:
             DefaultAzureCredential | AzureKeyCredential: The ai search credential
         """
-        if self.identity_type == IdentityType.SYSTEM_ASSIGNED:
+        if self.identity_type in [
+            IdentityType.SYSTEM_ASSIGNED,
+            IdentityType.USER_ASSIGNED,
+        ]:
             return DefaultAzureCredential()
-        elif self.identity_type == IdentityType.USER_ASSIGNED:
-            return DefaultAzureCredential(
-                managed_identity_client_id=self.ai_search_identity_id
-            )
         else:
             return AzureKeyCredential(
                 os.environ.get("AIService__AzureSearchOptions__Key")
@@ -180,7 +179,7 @@ def storage_account_blob_container_name(self) -> str:
         if container is None:
             raise ValueError(
                 f"""Populate environment variable 'StorageAccount__{
-                             self.normalised_indexer_type}__Container' with container name."""
+                    self.normalised_indexer_type}__Container' with container name."""
             )
 
         return container

text_2_sql/autogen/src/autogen_text_2_sql/creators/llm_model_creator.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from autogen_ext.models import AzureOpenAIChatCompletionClient
+from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
 from text_2_sql_core.utils.environment import IdentityType, get_identity_type
 
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
@@ -30,21 +30,15 @@ def get_model(cls, model_name: str) -> AzureOpenAIChatCompletionClient:
 
     @classmethod
     def get_authentication_properties(cls) -> dict:
-        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
+        if get_identity_type() in [
+            IdentityType.SYSTEM_ASSIGNED,
+            IdentityType.USER_ASSIGNED,
+        ]:
             # Create the token provider
             api_key = None
             token_provider = get_bearer_token_provider(
                 DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
             )
-        elif get_identity_type() == IdentityType.USER_ASSIGNED:
-            # Create the token provider
-            api_key = None
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["ClientId"]
-                ),
-                "https://cognitiveservices.azure.com/.default",
-            )
         else:
             token_provider = None
             api_key = os.environ["OpenAI__ApiKey"]
@@ -57,7 +51,7 @@ def gpt_4o_mini_model(cls) -> AzureOpenAIChatCompletionClient:
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__MiniCompletionDeployment"],
             model=os.environ["OpenAI__MiniCompletionDeployment"],
-            api_version="2024-08-01-preview",
+            api_version=os.environ["OpenAI__ApiVersion"],
             azure_endpoint=os.environ["OpenAI__Endpoint"],
             azure_ad_token_provider=token_provider,
             api_key=api_key,
@@ -75,7 +69,7 @@ def gpt_4o_model(cls) -> AzureOpenAIChatCompletionClient:
         return AzureOpenAIChatCompletionClient(
             azure_deployment=os.environ["OpenAI__CompletionDeployment"],
             model=os.environ["OpenAI__CompletionDeployment"],
-            api_version="2024-08-01-preview",
+            api_version=os.environ["OpenAI__ApiVersion"],
             azure_endpoint=os.environ["OpenAI__Endpoint"],
             azure_ad_token_provider=token_provider,
             api_key=api_key,

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/ai_search.py

@@ -1,9 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from azure.identity import DefaultAzureCredential
-from openai import AsyncAzureOpenAI
 from azure.core.credentials import AzureKeyCredential
-from azure.search.documents.models import VectorizedQuery, QueryType
+from azure.search.documents.models import QueryType, VectorizableTextQuery
 from azure.search.documents.aio import SearchClient
 from text_2_sql_core.utils.environment import IdentityType, get_identity_type
 import os
@@ -12,9 +11,13 @@
 from datetime import datetime, timezone
 import json
 from typing import Annotated
+from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
 
 class AISearchConnector:
+    def __init__(self):
+        self.open_ai_connector = OpenAIConnector()
+
     async def run_ai_search_query(
         self,
         query,
@@ -30,35 +33,18 @@ async def run_ai_search_query(
         identity_type = get_identity_type()
 
         if len(vector_fields) > 0:
-            async with AsyncAzureOpenAI(
-                # This is the default and can be omitted
-                api_key=os.environ["OpenAI__ApiKey"],
-                azure_endpoint=os.environ["OpenAI__Endpoint"],
-                api_version=os.environ["OpenAI__ApiVersion"],
-            ) as open_ai_client:
-                embeddings = await open_ai_client.embeddings.create(
-                    model=os.environ["OpenAI__EmbeddingModel"], input=query
-                )
-
-                # Extract the embedding vector
-                embedding_vector = embeddings.data[0].embedding
-
             vector_query = [
-                VectorizedQuery(
-                    vector=embedding_vector,
+                VectorizableTextQuery(
+                    text=query,
                     k_nearest_neighbors=7,
                     fields=",".join(vector_fields),
                 )
             ]
         else:
             vector_query = None
 
-        if identity_type == IdentityType.SYSTEM_ASSIGNED:
+        if identity_type in [IdentityType.SYSTEM_ASSIGNED, IdentityType.USER_ASSIGNED]:
             credential = DefaultAzureCredential()
-        elif identity_type == IdentityType.USER_ASSIGNED:
-            credential = DefaultAzureCredential(
-                managed_identity_client_id=os.environ["ClientID"]
-            )
         else:
             credential = AzureKeyCredential(
                 os.environ["AIService__AzureSearchOptions__Key"]
@@ -253,7 +239,9 @@ async def get_entity_schemas(
         logging.info("Filtered Schemas: %s", filtered_schemas)
         return filtered_schemas
 
-    async def add_entry_to_index(document: dict, vector_fields: dict, index_name: str):
+    async def add_entry_to_index(
+        self, document: dict, vector_fields: dict, index_name: str
+    ):
         """Add an entry to the search index."""
 
         logging.info("Document: %s", document)
@@ -270,31 +258,23 @@ async def add_entry_to_index(document: dict, vector_fields: dict, index_name: st
         document["DateLastModified"] = datetime.now(timezone.utc)
 
         try:
-            async with AsyncAzureOpenAI(
-                # This is the default and can be omitted
-                api_key=os.environ["OpenAI__ApiKey"],
-                azure_endpoint=os.environ["OpenAI__Endpoint"],
-                api_version=os.environ["OpenAI__ApiVersion"],
-            ) as open_ai_client:
-                embeddings = await open_ai_client.embeddings.create(
-                    model=os.environ["OpenAI__EmbeddingModel"],
-                    input=fields_to_embed.values(),
-                )
+            embeddings = await self.open_ai_connector.run_embedding_request(
+                list(fields_to_embed.values())
+            )
 
-                # Extract the embedding vector
-                for i, field in enumerate(vector_fields.values()):
-                    document[field] = embeddings.data[i].embedding
+            # Extract the embedding vector
+            for i, field in enumerate(vector_fields.values()):
+                document[field] = embeddings.data[i].embedding
 
             document["Id"] = base64.urlsafe_b64encode(
                 document["Question"].encode()
             ).decode("utf-8")
 
-            if identity_type == IdentityType.SYSTEM_ASSIGNED:
+            if identity_type in [
+                IdentityType.SYSTEM_ASSIGNED,
+                IdentityType.USER_ASSIGNED,
+            ]:
                 credential = DefaultAzureCredential()
-            elif identity_type == IdentityType.USER_ASSIGNED:
-                credential = DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["ClientID"]
-                )
             else:
                 credential = AzureKeyCredential(
                     os.environ["AIService__AzureSearchOptions__Key"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/open_ai.py

@@ -12,36 +12,61 @@
 class OpenAIConnector:
     @classmethod
     def get_authentication_properties(cls) -> dict:
-        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
+        if get_identity_type() in [
+            IdentityType.SYSTEM_ASSIGNED,
+            IdentityType.USER_ASSIGNED,
+        ]:
             # Create the token provider
             api_key = None
             token_provider = get_bearer_token_provider(
                 DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
             )
-        elif get_identity_type() == IdentityType.USER_ASSIGNED:
-            # Create the token provider
-            api_key = None
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["ClientId"]
-                ),
-                "https://cognitiveservices.azure.com/.default",
-            )
         else:
             token_provider = None
             api_key = os.environ["OpenAI__ApiKey"]
 
         return token_provider, api_key
 
-    async def run_completion_request(self, messages: list[dict], temperature=0):
+    async def run_completion_request(
+        self, messages: list[dict], temperature=0, max_tokens=2000, model="4o-mini"
+    ) -> str:
+        if model == "4o-mini":
+            model_deployment = os.environ["OpenAI__MiniCompletionDeployment"]
+        elif model == "4o":
+            model_deployment = os.environ["OpenAI__CompletionDeployment"]
+        else:
+            raise ValueError(f"Model {model} not found")
+
+        token_provider, api_key = self.get_authentication_properties()
         async with AsyncAzureOpenAI(
-            api_key=os.environ["OpenAI__ApiKey"],
-            azure_endpoint=os.environ["OpenAI__Endpoint"],
+            azure_deployment=model_deployment,
             api_version=os.environ["OpenAI__ApiVersion"],
+            azure_endpoint=os.environ["OpenAI__Endpoint"],
+            azure_ad_token_provider=token_provider,
+            api_key=api_key,
         ) as open_ai_client:
             response = await open_ai_client.chat.completions.create(
-                model=os.environ["OpenAI__MiniCompletionDeployment"],
+                model=model_deployment,
                 messages=messages,
                 temperature=temperature,
+                max_tokens=max_tokens,
             )
         return response.choices[0].message.content
+
+    async def run_embedding_request(self, batch: list[str]):
+        token_provider, api_key = self.get_authentication_properties()
+
+        model_deployment = os.environ["OpenAI__EmbeddingModel"]
+        async with AsyncAzureOpenAI(
+            azure_deployment=model_deployment,
+            api_version=os.environ["OpenAI__ApiVersion"],
+            azure_endpoint=os.environ["OpenAI__Endpoint"],
+            azure_ad_token_provider=token_provider,
+            api_key=api_key,
+        ) as open_ai_client:
+            embeddings = await open_ai_client.embeddings.create(
+                model=os.environ["OpenAI__EmbeddingModel"],
+                input=batch,
+            )
+
+        return embeddings

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/data_dictionary_creator.py

@@ -8,14 +8,12 @@
 import logging
 from pydantic import BaseModel, Field, ConfigDict, computed_field
 from typing import Optional
-from text_2_sql_core.utils.environment import IdentityType, get_identity_type
-from openai import AsyncAzureOpenAI
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import random
 import re
 import networkx as nx
 from text_2_sql_core.utils.database import DatabaseEngine
 from tenacity import retry, stop_after_attempt, wait_exponential
+from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
 logging.basicConfig(level=logging.INFO)
 
@@ -279,6 +277,8 @@ def __init__(
         if output_directory is None:
             self.output_directory = "."
 
+        self.open_ai_connector = OpenAIConnector()
+
         load_dotenv(find_dotenv())
 
     @property
@@ -627,9 +627,6 @@ async def extract_columns_with_definitions(
 
         return columns
 
-    @retry(
-        stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
-    )
     async def send_request_to_llm(self, system_prompt: str, input: str):
         """A method to use GPT to generate a definition for an entity.
 
@@ -640,60 +637,23 @@ async def send_request_to_llm(self, system_prompt: str, input: str):
         Returns:
             str: The generated definition."""
 
-        MAX_TOKENS = 2000
-
-        api_version = os.environ["OpenAI__ApiVersion"]
-        model = os.environ["OpenAI__CompletionDeployment"]
-
-        if get_identity_type() == IdentityType.SYSTEM_ASSIGNED:
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
-            )
-            api_key = None
-        elif get_identity_type() == IdentityType.USER_ASSIGNED:
-            token_provider = get_bearer_token_provider(
-                DefaultAzureCredential(
-                    managed_identity_client_id=os.environ["FunctionApp__ClientId"]
-                ),
-                "https://cognitiveservices.azure.com/.default",
-            )
-            api_key = None
-        else:
-            token_provider = None
-            api_key = os.environ["OpenAI__ApiKey"]
-
-        try:
-            async with AsyncAzureOpenAI(
-                api_key=api_key,
-                api_version=api_version,
-                azure_ad_token_provider=token_provider,
-                azure_endpoint=os.environ.get("OpenAI__Endpoint"),
-            ) as client:
-                response = await client.chat.completions.create(
-                    model=model,
-                    messages=[
-                        {
-                            "role": "system",
-                            "content": system_prompt,
-                        },
-                        {
-                            "role": "user",
-                            "content": [
-                                {
-                                    "type": "text",
-                                    "text": input,
-                                },
-                            ],
-                        },
-                    ],
-                    max_tokens=MAX_TOKENS,
-                )
+        messages = [
+            {
+                "role": "system",
+                "content": system_prompt,
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": input,
+                    },
+                ],
+            },
+        ]
 
-            return response.choices[0].message.content
-        except Exception as e:
-            logging.error(f"Unable to generate definition for {input}")
-            logging.error(f"Error generating definition: {e}")
-            return None
+        return await self.open_ai_connector.run_completion_request(messages)
 
     async def generate_entity_definition(self, entity: EntityItem):
         """A method to generate a definition for an entity.