From c379b9221aba85456bc539657a059ed2c1578663 Mon Sep 17 00:00:00 2001 From: Xu Yang Date: Wed, 23 Oct 2024 18:22:58 +0800 Subject: [PATCH] feat: independent use_azure_token_provider on chat and embedding (#452) * independent use_azure_token_provider * fix CI * fix doc format --- .env.example | 3 +- docs/installation_and_configuration.rst | 54 +++++++++++++------------ rdagent/oai/llm_conf.py | 3 +- rdagent/oai/llm_utils.py | 19 +++++---- 4 files changed, 44 insertions(+), 35 deletions(-) diff --git a/.env.example b/.env.example index 15976a3f..e43e8d0b 100644 --- a/.env.example +++ b/.env.example @@ -9,7 +9,8 @@ For more information about configuration options, please refer to the documentat # Global configs: USE_AZURE=False -USE_AZURE_TOKEN_PROVIDER=False +CHAT_USE_AZURE_TOKEN_PROVIDER=False +EMBEDDING_USE_AZURE_TOKEN_PROVIDER=False MAX_RETRY=10 RETRY_WAIT_SECONDS=20 diff --git a/docs/installation_and_configuration.rst b/docs/installation_and_configuration.rst index 2ac50d0d..536236d0 100644 --- a/docs/installation_and_configuration.rst +++ b/docs/installation_and_configuration.rst @@ -53,7 +53,7 @@ The following environment variables are standard configuration options for the u Use Azure Token Provider ------------------------ -If you are using the Azure token provider, you need to set the `USE_AZURE_TOKEN_PROVIDER` environment variable to `True`. then +If you are using the Azure token provider, you need to set the `CHAT_USE_AZURE_TOKEN_PROVIDER` and `EMBEDDING_USE_AZURE_TOKEN_PROVIDER` environment variable to `True`. then use the environment variables provided in the `Azure Configuration section `_. @@ -80,31 +80,33 @@ Configuration List - OpenAI API Setting -+-----------------------------+--------------------------------------------------+-------------------------+ -| Configuration Option | Meaning | Default Value | -+=============================+==================================================+=========================+ -| OPENAI_API_KEY | API key for both chat and embedding models | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| EMBEDDING_OPENAI_API_KEY | Use a different API key for embedding model | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| CHAT_OPENAI_API_KEY | Set to use a different API key for chat model | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| EMBEDDING_MODEL | Name of the embedding model | text-embedding-3-small | -+-----------------------------+--------------------------------------------------+-------------------------+ -| CHAT_MODEL | Name of the chat model | gpt-4-turbo | -+-----------------------------+--------------------------------------------------+-------------------------+ -| EMBEDDING_AZURE_API_BASE | Base URL for the Azure OpenAI API | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| EMBEDDING_AZURE_API_VERSION | Version of the Azure OpenAI API | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| CHAT_AZURE_API_BASE | Base URL for the Azure OpenAI API | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| CHAT_AZURE_API_VERSION | Version of the Azure OpenAI API | None | -+-----------------------------+--------------------------------------------------+-------------------------+ -| USE_AZURE | True if you are using Azure OpenAI | False | -+-----------------------------+--------------------------------------------------+-------------------------+ -| USE_AZURE_TOKEN_PROVIDER | True if you are using a Azure Token Provider | False | -+-----------------------------+--------------------------------------------------+-------------------------+ ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| Configuration Option | Meaning | Default Value | ++===================================+=================================================================+=========================+ +| OPENAI_API_KEY | API key for both chat and embedding models | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| EMBEDDING_OPENAI_API_KEY | Use a different API key for embedding model | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| CHAT_OPENAI_API_KEY | Set to use a different API key for chat model | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| EMBEDDING_MODEL | Name of the embedding model | text-embedding-3-small | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| CHAT_MODEL | Name of the chat model | gpt-4-turbo | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| EMBEDDING_AZURE_API_BASE | Base URL for the Azure OpenAI API | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| EMBEDDING_AZURE_API_VERSION | Version of the Azure OpenAI API | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| CHAT_AZURE_API_BASE | Base URL for the Azure OpenAI API | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| CHAT_AZURE_API_VERSION | Version of the Azure OpenAI API | None | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| USE_AZURE | True if you are using Azure OpenAI | False | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| CHAT_USE_AZURE_TOKEN_PROVIDER | True if you are using an Azure Token Provider in chat model | False | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ +| EMBEDDING_USE_AZURE_TOKEN_PROVIDER| True if you are using an Azure Token Provider in embedding model| False | ++-----------------------------------+-----------------------------------------------------------------+-------------------------+ - Globol Setting diff --git a/rdagent/oai/llm_conf.py b/rdagent/oai/llm_conf.py index ce40135d..09d9d8b5 100644 --- a/rdagent/oai/llm_conf.py +++ b/rdagent/oai/llm_conf.py @@ -9,7 +9,8 @@ class LLMSettings(BaseSettings): log_llm_chat_content: bool = True use_azure: bool = False - use_azure_token_provider: bool = False + chat_use_azure_token_provider: bool = False + embedding_use_azure_token_provider: bool = False managed_identity_client_id: str | None = None max_retry: int = 10 retry_wait_seconds: int = 1 diff --git a/rdagent/oai/llm_utils.py b/rdagent/oai/llm_utils.py index 921a5b5e..42b07b84 100644 --- a/rdagent/oai/llm_utils.py +++ b/rdagent/oai/llm_utils.py @@ -305,7 +305,8 @@ class APIBackend: self.encoder = None else: self.use_azure = LLM_SETTINGS.use_azure - self.use_azure_token_provider = LLM_SETTINGS.use_azure_token_provider + self.chat_use_azure_token_provider = LLM_SETTINGS.chat_use_azure_token_provider + self.embedding_use_azure_token_provider = LLM_SETTINGS.embedding_use_azure_token_provider self.managed_identity_client_id = LLM_SETTINGS.managed_identity_client_id # Priority: chat_api_key/embedding_api_key > openai_api_key > os.environ.get("OPENAI_API_KEY") @@ -341,7 +342,7 @@ class APIBackend: ) if self.use_azure: - if self.use_azure_token_provider: + if self.chat_use_azure_token_provider or self.embedding_use_azure_token_provider: dac_kwargs = {} if self.managed_identity_client_id is not None: dac_kwargs["managed_identity_client_id"] = self.managed_identity_client_id @@ -350,22 +351,26 @@ class APIBackend: credential, "https://cognitiveservices.azure.com/.default", ) + if self.chat_use_azure_token_provider: self.chat_client = openai.AzureOpenAI( azure_ad_token_provider=token_provider, api_version=self.chat_api_version, azure_endpoint=self.chat_api_base, ) - self.embedding_client = openai.AzureOpenAI( - azure_ad_token_provider=token_provider, - api_version=self.embedding_api_version, - azure_endpoint=self.embedding_api_base, - ) else: self.chat_client = openai.AzureOpenAI( api_key=self.chat_api_key, api_version=self.chat_api_version, azure_endpoint=self.chat_api_base, ) + + if self.embedding_use_azure_token_provider: + self.embedding_client = openai.AzureOpenAI( + azure_ad_token_provider=token_provider, + api_version=self.embedding_api_version, + azure_endpoint=self.embedding_api_base, + ) + else: self.embedding_client = openai.AzureOpenAI( api_key=self.embedding_api_key, api_version=self.embedding_api_version,