Add common embedding method

Author: BenConstable9

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,20 +258,13 @@ 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()

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

@@ -58,3 +58,21 @@ async def run_completion_request(
                 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