Skip to content

feat: independent use_azure_token_provider on chat and embedding #452

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
peteryang1 merged 3 commits into from
Oct 23, 2024
Merged

feat: independent use_azure_token_provider on chat and embedding #452

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9ad901e
independent use_azure_token_provider
peteryang1 Oct 23, 2024
16739ba
fix CI
peteryang1 Oct 23, 2024
42ce888
fix doc format
peteryang1 Oct 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
54 changes: 28 additions & 26 deletions docs/installation_and_configuration.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 <installation_and_configuration.html#azure-openai>`_.


Expand All @@ -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

Expand Down
3 changes: 2 additions & 1 deletion rdagent/oai/llm_conf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
19 changes: 12 additions & 7 deletions rdagent/oai/llm_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,7 +305,8 @@ def __init__( # noqa: C901, PLR0912, PLR0915
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")
Expand Down Expand Up @@ -341,7 +342,7 @@ def __init__( # noqa: C901, PLR0912, PLR0915
)

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
Expand All @@ -350,22 +351,26 @@ def __init__( # noqa: C901, PLR0912, PLR0915
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,
Expand Down