Migrate AzureOpenAI constructors to standard OpenAI client #2752
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,24 +2,21 @@ | |
| import asyncio | ||
| import logging | ||
| import os | ||
| from collections.abc import Awaitable, Callable | ||
| from enum import Enum | ||
| from typing import Optional | ||
|
|
||
| import aiohttp | ||
| from azure.core.credentials import AzureKeyCredential | ||
| from azure.core.credentials_async import AsyncTokenCredential | ||
| from azure.identity.aio import AzureDeveloperCliCredential, get_bearer_token_provider | ||
| from openai import AsyncOpenAI | ||
| from rich.logging import RichHandler | ||
|
|
||
| from load_azd_env import load_azd_env | ||
| from prepdocslib.blobmanager import BlobManager | ||
| from prepdocslib.csvparser import CsvParser | ||
| from prepdocslib.embeddings import ImageEmbeddings, OpenAIEmbeddings | ||
| from prepdocslib.fileprocessor import FileProcessor | ||
| from prepdocslib.filestrategy import FileStrategy | ||
| from prepdocslib.htmlparser import LocalHTMLParser | ||
|
|
@@ -160,17 +157,12 @@ class OpenAIHost(str, Enum): | |
|
|
||
|
|
||
| def setup_embeddings_service( | ||
| open_ai_client: AsyncOpenAI, | ||
| openai_host: OpenAIHost, | ||
| emb_model_name: str, | ||
| emb_model_dimensions: int, | ||
| azure_openai_deployment: str | None, | ||
| azure_openai_endpoint: str | None, | ||
| disable_vectors: bool = False, | ||
| disable_batch_vectors: bool = False, | ||
| ): | ||
|
|
@@ -179,70 +171,59 @@ def setup_embeddings_service( | |
| return None | ||
|
|
||
| if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]: | ||
| 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") | ||
|
|
||
| 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_openai_deployment, | ||
| azure_endpoint=azure_openai_endpoint, | ||
| ) | ||
|
|
||
|
|
||
| def setup_openai_client( | ||
| openai_host: OpenAIHost, | ||
| azure_credential: AsyncTokenCredential, | ||
| azure_openai_api_key: Optional[str] = None, | ||
| azure_openai_service: Optional[str] = None, | ||
| azure_openai_custom_url: Optional[str] = None, | ||
| openai_api_key: Optional[str] = None, | ||
| openai_organization: Optional[str] = None, | ||
| ) -> tuple[AsyncOpenAI, Optional[str]]: | ||
| openai_client: AsyncOpenAI | ||
| azure_openai_endpoint: Optional[str] = None | ||
|
|
||
| if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]: | ||
| base_url: Optional[str] = None | ||
| api_key_or_token: Optional[str | Callable[[], Awaitable[str]]] = None | ||
| if openai_host == OpenAIHost.AZURE_CUSTOM: | ||
| logger.info("OPENAI_HOST is azure_custom, setting up Azure OpenAI custom client") | ||
| if not azure_openai_custom_url: | ||
| raise ValueError("AZURE_OPENAI_CUSTOM_URL must be set when OPENAI_HOST is azure_custom") | ||
| base_url = azure_openai_custom_url | ||
| else: | ||
| 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") | ||
| azure_openai_endpoint = f"https://{azure_openai_service}.openai.azure.com" | ||
| base_url = f"{azure_openai_endpoint}/openai/v1" | ||
| if azure_openai_api_key: | ||
| logger.info("AZURE_OPENAI_API_KEY_OVERRIDE found, using as api_key for Azure OpenAI client") | ||
| api_key_or_token = azure_openai_api_key | ||
| else: | ||
| logger.info("Using Azure credential (passwordless authentication) for Azure OpenAI client") | ||
| api_key_or_token = get_bearer_token_provider( | ||
| azure_credential, "https://cognitiveservices.azure.com/.default" | ||
| ) | ||
| openai_client = AsyncOpenAI( | ||
| base_url=base_url, | ||
| api_key=api_key_or_token, # type: ignore[arg-type] | ||
| ) | ||
| elif openai_host == OpenAIHost.LOCAL: | ||
| logger.info("OPENAI_HOST is local, setting up local OpenAI client for OPENAI_BASE_URL with no key") | ||
| openai_client = AsyncOpenAI( | ||
|
|
@@ -259,7 +240,7 @@ def setup_openai_client( | |
| api_key=openai_api_key, | ||
| organization=openai_organization, | ||
| ) | ||
| return openai_client, azure_openai_endpoint | ||
|
|
||
|
|
||
| def setup_file_processors( | ||
|
|
@@ -368,7 +349,7 @@ async def main(strategy: Strategy, setup_index: bool = True): | |
| await strategy.run() | ||
|
|
||
|
|
||
| if __name__ == "__main__": # pragma: no cover | ||
| parser = argparse.ArgumentParser( | ||
| description="Prepare documents by extracting content from PDFs, splitting content into sections, uploading to blob storage, and indexing in a search index." | ||
| ) | ||
|
|
@@ -516,41 +497,34 @@ async def main(strategy: Strategy, setup_index: bool = True): | |
| enable_global_documents=enable_global_documents, | ||
| ) | ||
|
|
||
|
There was a problem hiding this comment. @copilot Isn't AZURE_OPENAI_API_VERSION in other files as well, like the GitHub actions workflows? There was a problem hiding this comment. You're absolutely right! Found and removed AZURE_OPENAI_API_VERSION references from GitHub Actions workflows (.github/workflows/azure-dev.yml, .github/workflows/evaluate.yaml), Azure DevOps pipeline (.azdo/pipelines/azure-dev.yml), and documentation (docs/reasoning.md). The only remaining reference is in evals/generate_ground_truth.py which uses LangChain wrappers, not the OpenAI SDK directly. Commit: 5d23055 |
||
| emb_model_dimensions = 1536 | ||
| if os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"): | ||
| emb_model_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"]) | ||
|
|
||
| openai_client, azure_openai_endpoint = setup_openai_client( | ||
| openai_host=OPENAI_HOST, | ||
| azure_credential=azd_credential, | ||
| azure_openai_service=os.getenv("AZURE_OPENAI_SERVICE"), | ||
| azure_openai_custom_url=os.getenv("AZURE_OPENAI_CUSTOM_URL"), | ||
| azure_openai_api_key=os.getenv("AZURE_OPENAI_API_KEY_OVERRIDE"), | ||
| openai_api_key=clean_key_if_exists(os.getenv("OPENAI_API_KEY")), | ||
| openai_organization=os.getenv("OPENAI_ORGANIZATION"), | ||
| ) | ||
| openai_embeddings_service = setup_embeddings_service( | ||
|
There was a problem hiding this comment. Moved after, since its now dependent on the client There was a problem hiding this comment. makes sense, and it's simpler at the end of the day |
||
| 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_openai_endpoint, | ||
| disable_vectors=dont_use_vectors, | ||
| disable_batch_vectors=args.disablebatchvectors, | ||
| ) | ||
|
|
||
| ingestion_strategy: Strategy | ||
| if use_int_vectorization: | ||
|
|
||
| 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( | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@copilot Remove this line that sets AZURE_OPENAI_ENDPOINT, and pass azure_openai_endpoint from
setup_openai_clientto setup_embeddings insteadThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed in commit ee14c51 - removed the redundant AZURE_OPENAI_ENDPOINT computation and now pass the
azure_openai_endpointvalue returned fromsetup_openai_clientdirectly tosetup_embeddings_service.