gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-08-07 20:17:45 +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
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,