gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-07-28 16:07:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: independent use_azure_token_provider on chat and embedding (#452)

Browse Source 
* independent use_azure_token_provider

* fix CI

* fix doc format
This commit is contained in:
Xu Yang
2024-10-23 18:22:58 +08:00
committed by GitHub
parent eace48a054
commit c71fd1fba2
4 changed files with 44 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+2 -1
View File
@@ -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
docs/installation_and_configuration.rst
+28 -26
View File
@@ -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>`_.
@@ -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
rdagent/oai/llm_conf.py
+2 -1
View File
@@ -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
rdagent/oai/llm_utils.py
+12 -7
View File
@@ -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,