feat/azure openai embedder (#265)

* add azure openai embedder

* bump changelog

* add int test

* directly map config in embedder process

* set default for api version

Author: rbiseck3

CHANGELOG.md

@@ -1,7 +1,8 @@
-## 0.3.4-dev1
+## 0.3.4
 
-### Features
+### Enhancements
 
+* **Add azure openai embedder**
 * **Add `collection_id` field to Couchbase `downloader_config`**
 
 ## 0.3.3

test/integration/embedders/test_azure_openai.py

@@ -0,0 +1,59 @@
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+from test.integration.embedders.utils import validate_embedding_output, validate_raw_embedder
+from test.integration.utils import requires_env
+from unstructured_ingest.embed.azure_openai import (
+    AzureOpenAIEmbeddingConfig,
+    AzureOpenAIEmbeddingEncoder,
+)
+from unstructured_ingest.v2.processes.embedder import Embedder, EmbedderConfig
+
+API_KEY = "AZURE_OPENAI_API_KEY"
+ENDPOINT = "AZURE_OPENAI_ENDPOINT"
+
+
+@dataclass(frozen=True)
+class AzureData:
+    api_key: str
+    endpoint: str
+
+
+def get_azure_data() -> AzureData:
+    api_key = os.getenv(API_KEY, None)
+    assert api_key
+    endpoint = os.getenv(ENDPOINT, None)
+    assert endpoint
+    return AzureData(api_key, endpoint)
+
+
+@requires_env(API_KEY, ENDPOINT)
+def test_azure_openai_embedder(embedder_file: Path):
+    azure_data = get_azure_data()
+    embedder_config = EmbedderConfig(
+        embedding_provider="azure-openai",
+        embedding_api_key=azure_data.api_key,
+        embedding_azure_endpoint=azure_data.endpoint,
+    )
+    embedder = Embedder(config=embedder_config)
+    results = embedder.run(elements_filepath=embedder_file)
+    assert results
+    with embedder_file.open("r") as f:
+        original_elements = json.load(f)
+    validate_embedding_output(original_elements=original_elements, output_elements=results)
+
+
+@requires_env(API_KEY, ENDPOINT)
+def test_raw_azure_openai_embedder(embedder_file: Path):
+    azure_data = get_azure_data()
+    embedder = AzureOpenAIEmbeddingEncoder(
+        config=AzureOpenAIEmbeddingConfig(
+            api_key=azure_data.api_key,
+            azure_endpoint=azure_data.endpoint,
+        )
+    )
+    validate_raw_embedder(
+        embedder=embedder, embedder_file=embedder_file, expected_dimensions=(1536,)
+    )

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.4-dev1"  # pragma: no cover
+__version__ = "0.3.4"  # pragma: no cover

unstructured_ingest/embed/azure_openai.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from pydantic import Field
+
+from unstructured_ingest.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from openai import AzureOpenAI
+
+
+class AzureOpenAIEmbeddingConfig(OpenAIEmbeddingConfig):
+    api_version: str = Field(description="Azure API version", default="2024-06-01")
+    azure_endpoint: str
+    embedder_model_name: str = Field(default="text-embedding-ada-002", alias="model_name")
+
+    @requires_dependencies(["openai"], extras="openai")
+    def get_client(self) -> "AzureOpenAI":
+        from openai import AzureOpenAI
+
+        return AzureOpenAI(
+            api_key=self.api_key.get_secret_value(),
+            api_version=self.api_version,
+            azure_endpoint=self.azure_endpoint,
+        )
+
+
+@dataclass
+class AzureOpenAIEmbeddingEncoder(OpenAIEmbeddingEncoder):
+    config: AzureOpenAIEmbeddingConfig

unstructured_ingest/v2/processes/embedder.py

@@ -16,6 +16,7 @@ class EmbedderConfig(BaseModel):
     embedding_provider: Optional[
         Literal[
             "openai",
+            "azure-openai",
             "huggingface",
             "aws-bedrock",
             "vertexai",
@@ -43,6 +44,14 @@ class EmbedderConfig(BaseModel):
     embedding_aws_region: Optional[str] = Field(
         default="us-west-2", description="AWS region used for AWS-based embedders, such as bedrock"
     )
+    embedding_azure_endpoint: Optional[str] = Field(
+        default=None,
+        description="Your Azure endpoint, including the resource, "
+        "e.g. `https://example-resource.azure.openai.com/`",
+    )
+    embedding_azure_api_version: Optional[str] = Field(
+        description="Azure API version", default=None
+    )
 
     def get_huggingface_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
         from unstructured_ingest.embed.huggingface import (
@@ -59,6 +68,25 @@ def get_openai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
 
         return OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig.model_validate(embedding_kwargs))
 
+    def get_azure_openai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
+        from unstructured_ingest.embed.azure_openai import (
+            AzureOpenAIEmbeddingConfig,
+            AzureOpenAIEmbeddingEncoder,
+        )
+
+        config_kwargs = {
+            "api_key": self.embedding_api_key,
+            "azure_endpoint": self.embedding_azure_endpoint,
+        }
+        if api_version := self.embedding_azure_api_version:
+            config_kwargs["api_version"] = api_version
+        if model_name := self.embedding_model_name:
+            config_kwargs["model_name"] = model_name
+
+        return AzureOpenAIEmbeddingEncoder(
+            config=AzureOpenAIEmbeddingConfig.model_validate(config_kwargs)
+        )
+
     def get_octoai_embedder(self, embedding_kwargs: dict) -> "BaseEmbeddingEncoder":
         from unstructured_ingest.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
 
@@ -146,6 +174,8 @@ def get_embedder(self) -> "BaseEmbeddingEncoder":
             return self.get_mixedbread_embedder(embedding_kwargs=kwargs)
         if self.embedding_provider == "togetherai":
             return self.get_togetherai_embedder(embedding_kwargs=kwargs)
+        if self.embedding_provider == "azure-openai":
+            return self.get_azure_openai_embedder(embedding_kwargs=kwargs)
 
         raise ValueError(f"{self.embedding_provider} not a recognized encoder")