Add openai setting

BenConstable9 · BenConstable9 · commit 97f32a6c028d · 2024-09-10T15:32:40.000+01:00
diff --git a/ai_search_with_adi/ai_search/.env b/ai_search_with_adi/ai_search/.env
@@ -0,0 +1,20 @@
+FunctionApp__Endpoint=<functionAppEndpoint>
+FunctionApp__Key=<functionAppKey>
+FunctionApp__PreEmbeddingCleaner__FunctionName=pre_embedding_cleaner
+FunctionApp__ADI__FunctionName=adi_2_ai_search
+FunctionApp__KeyPhraseExtractor__FunctionName=keyphrase_extractor
+FunctionApp__AppRegistrationResourceId=<App registration in form api://appRegistrationclientId if using identity based connections>
+AIService__AzureSearchOptions__IdentityType=<identityType> # system_assigned or user_assigned or key
+AIService__AzureSearchOptions__Endpoint=<searchServiceEndpoint>
+AIService__AzureSearchOptions__Identity__ClientId=<clientId if using user assigned identity>
+AIService__AzureSearchOptions__Key=<searchServiceKey if not using identity>
+AIService__AzureSearchOptions__UsePrivateEndpoint=<true/false>
+AIService__AzureSearchOptions__Identity__FQName=<fully qualified name of the identity if using user assigned identity>
+StorageAccount__FQEndpoint=<Fully qualified endpoint in form ResourceId=resourceId if using identity based connections>
+StorageAccount__ConnectionString=<connectionString if using non managed identity>
+StorageAccount__RagDocuments__Container=<containerName>
+OpenAI__ApiKey=<openAIKey if using non managed identity>
+OpenAI__Endpoint=<openAIEndpoint>
+OpenAI__EmbeddingModel=<openAIEmbeddingModelName>
+OpenAI__EmbeddingDeployment=<openAIEmbeddingDeploymentId>
+OpenAI__EmbeddingDimensions=1536
diff --git a/ai_search_with_adi/ai_search/ai_search.py b/ai_search_with_adi/ai_search/ai_search.py
@@ -307,13 +307,13 @@ def get_adi_skill(self, chunk_by_page=False) -> WebApiSkill:
         )
 
         if self.environment.identity_type != IdentityType.KEY:
-            adi_skill.auth_identity = self.environment.function_app_app_registration_resource_id
-
-        if self.environment.identity_type == IdentityType.USER_ASSIGNED:
             adi_skill.auth_identity = (
-                self.environment.ai_search_user_assigned_identity
+                self.environment.function_app_app_registration_resource_id
             )
 
+        if self.environment.identity_type == IdentityType.USER_ASSIGNED:
+            adi_skill.auth_identity = self.environment.ai_search_user_assigned_identity
+
         return adi_skill
 
     def get_vector_skill(
@@ -335,12 +335,20 @@ def get_vector_skill(
             name="Vector Skill",
             description="Skill to generate embeddings",
             context=context,
-            deployment_id="0",
-            model_name="text-embedding-3-large",
+            deployment_id=self.environment.open_ai_embedding_deployment,
+            model_name=self.environment.open_ai_embedding_model,
             inputs=embedding_skill_inputs,
             outputs=embedding_skill_outputs,
+            dimensions=self.environment.open_ai_embedding_dimensions,
         )
 
+        if self.environment.identity_type == IdentityType.KEY:
+            vector_skill.api_key = self.environment.open_ai_api_key
+        elif self.environment.identity_type == IdentityType.USER_ASSIGNED:
+            vector_skill.auth_identity = (
+                self.environment.ai_search_user_assigned_identity
+            )
+
         return vector_skill
 
     def get_key_phrase_extraction_skill(self, context, source) -> WebApiSkill:
@@ -403,6 +411,19 @@ def get_vector_search(self) -> VectorSearch:
             VectorSearch: The vector search configuration
         """
 
+        open_ai_params = AzureOpenAIParameters(
+            resource_uri=self.environment.open_ai_endpoint,
+            modelName=self.environment.open_ai_embedding_model,
+            deploymentId=self.environment.open_ai_embedding_deployment,
+        )
+
+        if self.environment.identity_type == IdentityType.KEY:
+            open_ai_params.api_key = self.environment.open_ai_api_key
+        elif self.environment.identity_type == IdentityType.USER_ASSIGNED:
+            open_ai_params.auth_identity = (
+                self.environment.ai_search_user_assigned_identity
+            )
+
         vector_search = VectorSearch(
             algorithms=[
                 HnswAlgorithmConfiguration(name=self.algorithm_name),
@@ -417,7 +438,7 @@ def get_vector_search(self) -> VectorSearch:
             vectorizers=[
                 AzureOpenAIVectorizer(
                     name=self.vectorizer_name,
-                    azure_open_ai_parameters=AzureOpenAIParameters(),
+                    azure_open_ai_parameters=open_ai_params,
                 ),
             ],
         )
diff --git a/ai_search_with_adi/ai_search/environment.py b/ai_search_with_adi/ai_search/environment.py
@@ -7,20 +7,24 @@
 from azure.core.credentials import AzureKeyCredential
 from azure.search.documents.indexes.models import SearchIndexerDataUserAssignedIdentity
 
+
 class IndexerType(Enum):
     """The type of the indexer"""
 
     RAG_DOCUMENTS = "rag-documents"
 
+
 class IdentityType(Enum):
     """The type of the indexer"""
 
     USER_ASSIGNED = "user_assigned"
     SYSTEM_ASSIGNED = "system_assigned"
     KEY = "key"
 
+
 class AISearchEnvironment:
     """This class is used to get the environment variables for the AI search service."""
+
     def __init__(self, indexer_type: IndexerType):
         """Initialize the AISearchEnvironment class.
 
@@ -33,7 +37,7 @@ def __init__(self, indexer_type: IndexerType):
     @property
     def normalised_indexer_type(self) -> str:
         """This function returns the normalised indexer type.
-        
+
         Returns:
             str: The normalised indexer type
         """
@@ -46,7 +50,7 @@ def normalised_indexer_type(self) -> str:
     @property
     def identity_type(self) -> IdentityType:
         """This function returns the identity type.
-        
+
         Returns:
             IdentityType: The identity type
         """
@@ -60,54 +64,100 @@ def identity_type(self) -> IdentityType:
             return IdentityType.KEY
         else:
             raise ValueError("Invalid identity type")
-        
+
     @property
     def ai_search_endpoint(self) -> str:
         """This function returns the ai search endpoint.
-        
+
         Returns:
             str: The ai search endpoint
         """
         return os.environ.get("AIService__AzureSearchOptions__Endpoint")
-    
+
     @property
     def ai_search_identity_id(self) -> str:
         """This function returns the ai search identity id.
-        
+
         Returns:
             str: The ai search identity id
         """
         return os.environ.get("AIService__AzureSearchOptions__Identity__ClientId")
-    
+
     @property
     def ai_search_user_assigned_identity(self) -> SearchIndexerDataUserAssignedIdentity:
         """This function returns the ai search user assigned identity.
-        
+
         Returns:
-            SearchIndexerDataUserAssignedIdentity: The ai search user assigned identity"""
+            SearchIndexerDataUserAssignedIdentity: The ai search user assigned identity
+        """
         user_assigned_identity = SearchIndexerDataUserAssignedIdentity(
-            user_assigned_identity=os.environ.get("AIService__AzureSearchOptions__Identity__FQName")
+            user_assigned_identity=os.environ.get(
+                "AIService__AzureSearchOptions__Identity__FQName"
+            )
         )
         return user_assigned_identity
 
     @property
     def ai_search_credential(self) -> DefaultAzureCredential | AzureKeyCredential:
         """This function returns the ai search credential.
-        
+
         Returns:
             DefaultAzureCredential | AzureKeyCredential: The ai search credential
         """
         if self.identity_type in IdentityType.SYSTEM_ASSIGNED:
             return DefaultAzureCredential()
         elif self.identity_type in IdentityType.USER_ASSIGNED:
-            return DefaultAzureCredential(managed_identity_client_id=self.ai_search_identity_id)
+            return DefaultAzureCredential(
+                managed_identity_client_id=self.ai_search_identity_id
+            )
         else:
-            return AzureKeyCredential(os.environ.get("AIService__AzureSearchOptions__Key"))
+            return AzureKeyCredential(
+                os.environ.get("AIService__AzureSearchOptions__Key")
+            )
+
+    @property
+    def open_ai_api_key(self) -> str:
+        """This function returns the open ai api key.
+
+        Returns:
+            str: The open ai api key
+        """
+        return os.environ.get("OpenAI__ApiKey")
+
+    @property
+    def open_ai_endpoint(self) -> str:
+        """This function returns the open ai endpoint.
+
+        Returns:
+            str: The open ai endpoint
+        """
+        return os.environ.get("OpenAI__Endpoint")
+
+    @property
+    def open_ai_embedding_model(self) -> str:
+        """This function returns the open ai embedding model.
+
+        Returns:
+            str: The open ai embedding model
+        """
+        return os.environ.get("OpenAI__EmbeddingModel")
+
+    @property
+    def open_ai_embedding_deployment(self) -> str:
+        """This function returns the open ai embedding deployment.
+
+        Returns:
+            str: The open ai embedding deployment
+        """
+        return os.environ.get("OpenAI__EmbeddingDeployment")
 
     @property
     def storage_account_connection_string(self) -> str:
         """This function returns the blob connection string. If the identity type is user_assigned or system_assigned, it returns the FQEndpoint, otherwise it returns the ConnectionString"""
-        if self.identity_type in [IdentityType.SYSTEM_ASSIGNED, IdentityType.USER_ASSIGNED]:
+        if self.identity_type in [
+            IdentityType.SYSTEM_ASSIGNED,
+            IdentityType.USER_ASSIGNED,
+        ]:
             return os.environ.get("StorageAccount__FQEndpoint")
         else:
             return os.environ.get("StorageAccount__ConnectionString")
@@ -118,8 +168,10 @@ def storage_account_blob_container_name(self) -> str:
         This function returns azure blob container name
         """
 
-        return os.environ.get(f"StorageAccount__{self.normalised_indexer_type}__Container")
-    
+        return os.environ.get(
+            f"StorageAccount__{self.normalised_indexer_type}__Container"
+        )
+
     @property
     def function_app_end_point(self) -> str:
         """
@@ -133,14 +185,14 @@ def function_app_key(self) -> str:
         This function returns function app key
         """
         return os.environ.get("FunctionApp__Key")
-    
+
     @property
     def function_app_app_registration_resource_id(self) -> str:
         """
         This function returns function app app registration resource id
         """
         return os.environ.get("FunctionApp__AppRegistrationResourceId")
-    
+
     @property
     def function_app_pre_embedding_cleaner_route(self) -> str:
         """
@@ -161,26 +213,27 @@ def function_app_key_phrase_extractor_route(self) -> str:
         This function returns function app keyphrase extractor name
         """
         return os.environ.get("FunctionApp__KeyPhraseExtractor__FunctionName")
-    
+
     @property
-    def ai_search_embedding_model_dimensions(self) -> str:
+    def open_ai_embedding_dimensions(self) -> str:
         """
         This function returns dimensions for embedding model.
 
         Returns:
             str: The dimensions for embedding model
         """
 
-        return os.environ.get(
-            f"AIService__AzureSearchOptions__{self.normalised_indexer_type}__EmbeddingDimensions"
-        )
-    
+        return os.environ.get("OpenAI__EmbeddingDimensions")
+
     @property
     def use_private_endpoint(self) -> bool:
         """
         This function returns true if private endpoint is used
         """
-        return os.environ.get("AIService__AzureSearchOptions__UsePrivateEndpoint") == "true"
+        return (
+            os.environ.get("AIService__AzureSearchOptions__UsePrivateEndpoint")
+            == "true"
+        )
 
     def get_custom_skill_function_url(self, skill_type: str):
         """
@@ -194,7 +247,9 @@ def get_custom_skill_function_url(self, skill_type: str):
             route = self.function_app_key_phrase_extractor_route
         else:
             raise ValueError(f"Invalid skill type: {skill_type}")
-        
-        full_url = f"{self.function_app_end_point}/api/{route}?code={self.function_app_key}"
 
-        return full_url
+        full_url = (
+            f"{self.function_app_end_point}/api/{route}?code={self.function_app_key}"
+        )
+
+        return full_url
diff --git a/ai_search_with_adi/ai_search/rag_documents.py b/ai_search_with_adi/ai_search/rag_documents.py
@@ -83,7 +83,7 @@ def get_index_fields(self) -> list[SearchableField]:
             SearchField(
                 name="ChunkEmbedding",
                 type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
-                vector_search_dimensions=self.environment.ai_search_embedding_model_dimensions,
+                vector_search_dimensions=self.environment.open_ai_embedding_dimensions,
                 vector_search_profile_name=self.vector_search_profile_name,
             ),
             SearchableField(
diff --git a/ai_search_with_adi/function_apps/indexer/pre_embedding_cleaner.py b/ai_search_with_adi/function_apps/indexer/pre_embedding_cleaner.py
