Refactoring of embeddings to use the same client

Author: pamelafox

app/backend/app.py

@@ -90,7 +90,6 @@
 from error import error_dict, error_response
 from prepdocs import (
     OpenAIHost,
-    clean_key_if_exists,
     setup_embeddings_service,
     setup_file_processors,
     setup_image_embeddings_service,
@@ -426,6 +425,11 @@ async def setup_clients():
         os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT") if OPENAI_HOST in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM] else None
     )
     AZURE_OPENAI_CUSTOM_URL = os.getenv("AZURE_OPENAI_CUSTOM_URL")
+    AZURE_OPENAI_ENDPOINT = (
+        os.getenv("AZURE_OPENAI_ENDPOINT")
+        or (AZURE_OPENAI_CUSTOM_URL if OPENAI_HOST == OpenAIHost.AZURE_CUSTOM else None)
+        or (f"https://{AZURE_OPENAI_SERVICE}.openai.azure.com" if AZURE_OPENAI_SERVICE else None)
+    )
     AZURE_VISION_ENDPOINT = os.getenv("AZURE_VISION_ENDPOINT", "")
     AZURE_OPENAI_API_KEY_OVERRIDE = os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE")
     # Used only with non-Azure OpenAI deployments
@@ -599,16 +603,12 @@ async def setup_clients():
             search_service=AZURE_SEARCH_SERVICE, index_name=AZURE_SEARCH_INDEX, azure_credential=azure_credential
         )
         text_embeddings_service = setup_embeddings_service(
-            azure_credential=azure_credential,
-            openai_host=OpenAIHost(OPENAI_HOST),
+            open_ai_client=openai_client,
+            openai_host=OPENAI_HOST,
             emb_model_name=OPENAI_EMB_MODEL,
             emb_model_dimensions=OPENAI_EMB_DIMENSIONS,
-            azure_openai_service=AZURE_OPENAI_SERVICE,
-            azure_openai_custom_url=AZURE_OPENAI_CUSTOM_URL,
             azure_openai_deployment=AZURE_OPENAI_EMB_DEPLOYMENT,
-            azure_openai_key=clean_key_if_exists(AZURE_OPENAI_API_KEY_OVERRIDE),
-            openai_key=clean_key_if_exists(OPENAI_API_KEY),
-            openai_org=OPENAI_ORGANIZATION,
+            azure_openai_endpoint=AZURE_OPENAI_ENDPOINT,
             disable_vectors=os.getenv("USE_VECTORS", "").lower() == "false",
         )
         image_embeddings_service = setup_image_embeddings_service(

app/backend/prepdocs.py

@@ -15,11 +15,7 @@
 from load_azd_env import load_azd_env
 from prepdocslib.blobmanager import BlobManager
 from prepdocslib.csvparser import CsvParser
-from prepdocslib.embeddings import (
-    AzureOpenAIEmbeddingService,
-    ImageEmbeddings,
-    OpenAIEmbeddingService,
-)
+from prepdocslib.embeddings import ImageEmbeddings, OpenAIEmbeddings
 from prepdocslib.fileprocessor import FileProcessor
 from prepdocslib.filestrategy import FileStrategy
 from prepdocslib.htmlparser import LocalHTMLParser
@@ -160,46 +156,37 @@ class OpenAIHost(str, Enum):
 
 
 def setup_embeddings_service(
-    azure_credential: AsyncTokenCredential,
+    open_ai_client: AsyncOpenAI,
     openai_host: OpenAIHost,
     emb_model_name: str,
     emb_model_dimensions: int,
-    azure_openai_service: Optional[str],
-    azure_openai_custom_url: Optional[str],
-    azure_openai_deployment: Optional[str],
-    azure_openai_key: Optional[str],
-    openai_key: Optional[str],
-    openai_org: Optional[str],
+    azure_openai_deployment: str | None,
+    azure_openai_endpoint: str | None,
     disable_vectors: bool = False,
     disable_batch_vectors: bool = False,
 ):
     if disable_vectors:
         logger.info("Not setting up embeddings service")
         return None
 
+    azure_endpoint = None
+    azure_deployment = None
     if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]:
-        azure_open_ai_credential: AsyncTokenCredential | AzureKeyCredential = (
-            azure_credential if azure_openai_key is None else AzureKeyCredential(azure_openai_key)
-        )
-        return AzureOpenAIEmbeddingService(
-            open_ai_service=azure_openai_service,
-            open_ai_custom_url=azure_openai_custom_url,
-            open_ai_deployment=azure_openai_deployment,
-            open_ai_model_name=emb_model_name,
-            open_ai_dimensions=emb_model_dimensions,
-            credential=azure_open_ai_credential,
-            disable_batch=disable_batch_vectors,
-        )
-    else:
-        if openai_key is None:
-            raise ValueError("OpenAI key is required when using the non-Azure OpenAI API")
-        return OpenAIEmbeddingService(
-            open_ai_model_name=emb_model_name,
-            open_ai_dimensions=emb_model_dimensions,
-            credential=openai_key,
-            organization=openai_org,
-            disable_batch=disable_batch_vectors,
-        )
+        if azure_openai_endpoint is None:
+            raise ValueError("Azure OpenAI endpoint must be provided when using Azure OpenAI embeddings")
+        if azure_openai_deployment is None:
+            raise ValueError("Azure OpenAI deployment must be provided when using Azure OpenAI embeddings")
+        azure_endpoint = azure_openai_endpoint
+        azure_deployment = azure_openai_deployment
+
+    return OpenAIEmbeddings(
+        open_ai_client=open_ai_client,
+        open_ai_model_name=emb_model_name,
+        open_ai_dimensions=emb_model_dimensions,
+        disable_batch=disable_batch_vectors,
+        azure_deployment_name=azure_deployment,
+        azure_endpoint=azure_endpoint,
+    )
 
 
 def setup_openai_client(
@@ -226,17 +213,15 @@ def setup_openai_client(
             logger.info("OPENAI_HOST is azure, setting up Azure OpenAI client")
             if not azure_openai_service:
                 raise ValueError("AZURE_OPENAI_SERVICE must be set when OPENAI_HOST is azure")
-            endpoint = f"https://{azure_openai_service}.openai.azure.com"
+            endpoint = f"https://{azure_openai_service}.openai.azure.com/openai/v1"
         if azure_openai_api_key:
             logger.info("AZURE_OPENAI_API_KEY_OVERRIDE found, using as api_key for Azure OpenAI client")
-            openai_client = AsyncOpenAI(
-                base_url=f"{endpoint}/openai/v1", api_key=azure_openai_api_key
-            )
+            openai_client = AsyncOpenAI(base_url=endpoint, api_key=azure_openai_api_key)
         else:
             logger.info("Using Azure credential (passwordless authentication) for Azure OpenAI client")
             token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default")
             openai_client = AsyncOpenAI(
-                base_url=f"{endpoint}/openai/v1",
+                base_url=endpoint,
                 api_key=token_provider,
             )
     elif openai_host == OpenAIHost.LOCAL:
@@ -515,20 +500,6 @@ async def main(strategy: Strategy, setup_index: bool = True):
     emb_model_dimensions = 1536
     if os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"):
         emb_model_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"])
-    openai_embeddings_service = setup_embeddings_service(
-        azure_credential=azd_credential,
-        openai_host=OPENAI_HOST,
-        emb_model_name=os.environ["AZURE_OPENAI_EMB_MODEL_NAME"],
-        emb_model_dimensions=emb_model_dimensions,
-        azure_openai_service=os.getenv("AZURE_OPENAI_SERVICE"),
-        azure_openai_custom_url=os.getenv("AZURE_OPENAI_CUSTOM_URL"),
-        azure_openai_deployment=os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT"),
-        azure_openai_key=os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE"),
-        openai_key=clean_key_if_exists(os.getenv("OPENAI_API_KEY")),
-        openai_org=os.getenv("OPENAI_ORGANIZATION"),
-        disable_vectors=dont_use_vectors,
-        disable_batch_vectors=args.disablebatchvectors,
-    )
     openai_client = setup_openai_client(
         openai_host=OPENAI_HOST,
         azure_credential=azd_credential,
@@ -538,11 +509,25 @@ async def main(strategy: Strategy, setup_index: bool = True):
         openai_api_key=clean_key_if_exists(os.getenv("OPENAI_API_KEY")),
         openai_organization=os.getenv("OPENAI_ORGANIZATION"),
     )
+    azure_embedding_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") or os.getenv("AZURE_OPENAI_CUSTOM_URL")
+    if not azure_embedding_endpoint and OPENAI_HOST == OpenAIHost.AZURE:
+        if service := os.getenv("AZURE_OPENAI_SERVICE"):
+            azure_embedding_endpoint = f"https://{service}.openai.azure.com"
+    openai_embeddings_service = setup_embeddings_service(
+        open_ai_client=openai_client,
+        openai_host=OPENAI_HOST,
+        emb_model_name=os.environ["AZURE_OPENAI_EMB_MODEL_NAME"],
+        emb_model_dimensions=emb_model_dimensions,
+        azure_openai_deployment=os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT"),
+        azure_openai_endpoint=azure_embedding_endpoint,
+        disable_vectors=dont_use_vectors,
+        disable_batch_vectors=args.disablebatchvectors,
+    )
 
     ingestion_strategy: Strategy
     if use_int_vectorization:
 
-        if not openai_embeddings_service or not isinstance(openai_embeddings_service, AzureOpenAIEmbeddingService):
+        if not openai_embeddings_service or OPENAI_HOST not in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]:
             raise Exception("Integrated vectorization strategy requires an Azure OpenAI embeddings service")
 
         ingestion_strategy = IntegratedVectorizerStrategy(

app/backend/prepdocslib/embeddings.py

@@ -1,14 +1,10 @@
 import logging
 from abc import ABC
 from collections.abc import Awaitable, Callable
-from typing import Optional
 from urllib.parse import urljoin
 
 import aiohttp
 import tiktoken
-from azure.core.credentials import AzureKeyCredential
-from azure.core.credentials_async import AsyncTokenCredential
-from azure.identity.aio import get_bearer_token_provider
 from openai import AsyncOpenAI, RateLimitError
 from tenacity import (
     AsyncRetrying,
@@ -22,9 +18,7 @@
 
 
 class EmbeddingBatch:
-    """
-    Represents a batch of text that is going to be embedded
-    """
+    """Represents a batch of text that is going to be embedded."""
 
     def __init__(self, texts: list[str], token_length: int):
         self.texts = texts
@@ -36,12 +30,9 @@ class ExtraArgs(TypedDict, total=False):
 
 
 class OpenAIEmbeddings(ABC):
-    """
-    Contains common logic across both OpenAI and Azure OpenAI embedding services
-    Can split source text into batches for more efficient embedding calls
-    """
+    """Client wrapper that handles batching, retries, and token accounting."""
 
-    SUPPORTED_BATCH_AOAI_MODEL = {
+    SUPPORTED_BATCH_MODEL = {
         "text-embedding-ada-002": {"token_limit": 8100, "max_batch_size": 16},
         "text-embedding-3-small": {"token_limit": 8100, "max_batch_size": 16},
         "text-embedding-3-large": {"token_limit": 8100, "max_batch_size": 16},
@@ -52,13 +43,26 @@ class OpenAIEmbeddings(ABC):
         "text-embedding-3-large": True,
     }
 
-    def __init__(self, open_ai_model_name: str, open_ai_dimensions: int, disable_batch: bool = False):
+    def __init__(
+        self,
+        open_ai_client: AsyncOpenAI,
+        open_ai_model_name: str,
+        open_ai_dimensions: int,
+        *,
+        disable_batch: bool = False,
+        azure_deployment_name: str | None = None,
+        azure_endpoint: str | None = None,
+    ):
+        self.open_ai_client = open_ai_client
         self.open_ai_model_name = open_ai_model_name
         self.open_ai_dimensions = open_ai_dimensions
         self.disable_batch = disable_batch
+        self.azure_deployment_name = azure_deployment_name
+        self.azure_endpoint = azure_endpoint.rstrip("/") if azure_endpoint else None
 
-    async def create_client(self) -> AsyncOpenAI:
-        raise NotImplementedError
+    @property
+    def _api_model(self) -> str:
+        return self.azure_deployment_name or self.open_ai_model_name
 
     def before_retry_sleep(self, retry_state):
         logger.info("Rate limited on the OpenAI embeddings API, sleeping before retrying...")
@@ -68,7 +72,7 @@ def calculate_token_length(self, text: str):
         return len(encoding.encode(text))
 
     def split_text_into_batches(self, texts: list[str]) -> list[EmbeddingBatch]:
-        batch_info = OpenAIEmbeddings.SUPPORTED_BATCH_AOAI_MODEL.get(self.open_ai_model_name)
+        batch_info = OpenAIEmbeddings.SUPPORTED_BATCH_MODEL.get(self.open_ai_model_name)
         if not batch_info:
             raise NotImplementedError(
                 f"Model {self.open_ai_model_name} is not supported with batch embedding operations"
@@ -101,7 +105,6 @@ def split_text_into_batches(self, texts: list[str]) -> list[EmbeddingBatch]:
     async def create_embedding_batch(self, texts: list[str], dimensions_args: ExtraArgs) -> list[list[float]]:
         batches = self.split_text_into_batches(texts)
         embeddings = []
-        client = await self.create_client()
         for batch in batches:
             async for attempt in AsyncRetrying(
                 retry=retry_if_exception_type(RateLimitError),
@@ -110,8 +113,8 @@ async def create_embedding_batch(self, texts: list[str], dimensions_args: ExtraA
                 before_sleep=self.before_retry_sleep,
             ):
                 with attempt:
-                    emb_response = await client.embeddings.create(
-                        model=self.open_ai_model_name, input=batch.texts, **dimensions_args
+                    emb_response = await self.open_ai_client.embeddings.create(
+                        model=self._api_model, input=batch.texts, **dimensions_args
                     )
                     embeddings.extend([data.embedding for data in emb_response.data])
                     logger.info(
@@ -123,16 +126,15 @@ async def create_embedding_batch(self, texts: list[str], dimensions_args: ExtraA
         return embeddings
 
     async def create_embedding_single(self, text: str, dimensions_args: ExtraArgs) -> list[float]:
-        client = await self.create_client()
         async for attempt in AsyncRetrying(
             retry=retry_if_exception_type(RateLimitError),
             wait=wait_random_exponential(min=15, max=60),
             stop=stop_after_attempt(15),
             before_sleep=self.before_retry_sleep,
         ):
             with attempt:
-                emb_response = await client.embeddings.create(
-                    model=self.open_ai_model_name, input=text, **dimensions_args
+                emb_response = await self.open_ai_client.embeddings.create(
+                    model=self._api_model, input=text, **dimensions_args
                 )
                 logger.info("Computed embedding for text section. Character count: %d", len(text))
 
@@ -146,85 +148,12 @@ async def create_embeddings(self, texts: list[str]) -> list[list[float]]:
             else {}
         )
 
-        if not self.disable_batch and self.open_ai_model_name in OpenAIEmbeddings.SUPPORTED_BATCH_AOAI_MODEL:
+        if not self.disable_batch and self.open_ai_model_name in OpenAIEmbeddings.SUPPORTED_BATCH_MODEL:
             return await self.create_embedding_batch(texts, dimensions_args)
 
         return [await self.create_embedding_single(text, dimensions_args) for text in texts]
 
 
-class AzureOpenAIEmbeddingService(OpenAIEmbeddings):
-    """
-    Class for using Azure OpenAI embeddings
-    To learn more please visit https://learn.microsoft.com/azure/ai-services/openai/concepts/understand-embeddings
-    """
-
-    def __init__(
-        self,
-        open_ai_service: Optional[str],
-        open_ai_deployment: Optional[str],
-        open_ai_model_name: str,
-        open_ai_dimensions: int,
-        credential: AsyncTokenCredential | AzureKeyCredential,
-        open_ai_custom_url: Optional[str] = None,
-        disable_batch: bool = False,
-    ):
-        super().__init__(open_ai_deployment or open_ai_model_name, open_ai_dimensions, disable_batch)
-        self.open_ai_service = open_ai_service
-        if open_ai_service:
-            self.open_ai_endpoint = f"https://{open_ai_service}.openai.azure.com"
-        elif open_ai_custom_url:
-            self.open_ai_endpoint = open_ai_custom_url
-        else:
-            raise ValueError("Either open_ai_service or open_ai_custom_url must be provided")
-        self.open_ai_deployment = open_ai_deployment
-        self.credential = credential
-
-    async def create_client(self) -> AsyncOpenAI:
-        class AuthArgs(TypedDict, total=False):
-            api_key: str
-
-        auth_args = AuthArgs()
-        if isinstance(self.credential, AzureKeyCredential):
-            auth_args["api_key"] = self.credential.key
-        elif isinstance(self.credential, AsyncTokenCredential):
-            token_provider = get_bearer_token_provider(
-                self.credential, "https://cognitiveservices.azure.com/.default"
-            )
-            auth_args["api_key"] = token_provider
-        else:
-            raise TypeError("Invalid credential type")
-
-        # For Azure OpenAI, we need to use the v1 endpoint
-        base_url = f"{self.open_ai_endpoint}/openai/v1"
-        
-        return AsyncOpenAI(
-            base_url=base_url,
-            **auth_args,
-        )
-
-
-class OpenAIEmbeddingService(OpenAIEmbeddings):
-    """
-    Class for using OpenAI embeddings
-    To learn more please visit https://platform.openai.com/docs/guides/embeddings
-    """
-
-    def __init__(
-        self,
-        open_ai_model_name: str,
-        open_ai_dimensions: int,
-        credential: str,
-        organization: Optional[str] = None,
-        disable_batch: bool = False,
-    ):
-        super().__init__(open_ai_model_name, open_ai_dimensions, disable_batch)
-        self.credential = credential
-        self.organization = organization
-
-    async def create_client(self) -> AsyncOpenAI:
-        return AsyncOpenAI(api_key=self.credential, organization=self.organization)
-
-
 class ImageEmbeddings:
     """
     Class for using image embeddings from Azure AI Vision