yujunjun/autogen
1
0
Fork 0
mirror of https://github.com/microsoft/autogen.git synced 2025-09-02 12:57:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added a flag to agbench to enable Azure identity. (#5977)

Browse Source
This commit is contained in:
afourney 2025-03-17 00:10:44 -07:00 committed by GitHub
parent 483532180a
commit 22b68b96b6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
python/packages/agbench/src/agbench/run_cmd.py
View File

@ -284,14 +284,8 @@ def get_scenario_env(token_provider: Optional[Callable[[], str]] = None, env_fil
## Support Azure auth tokens
azure_openai_ad_token = os.environ.get("AZURE_OPENAI_AD_TOKEN")
if not azure_openai_ad_token and token_provider:
if azure_openai_ad_token is None and token_provider is not None:
azure_openai_ad_token = token_provider()
if not azure_openai_ad_token:
azure_token_provider = get_azure_token_provider()
if azure_token_provider:
azure_openai_ad_token = azure_token_provider()
else:
logging.warning("No Azure AD token provider found. Azure AD token not set.")
if azure_openai_ad_token is not None and len(azure_openai_ad_token.strip()) > 0:
env["AZURE_OPENAI_AD_TOKEN"] = azure_openai_ad_token
@ -888,6 +882,12 @@ def run_cli(args: Sequence[str]) -> None:
help="The number of parallel processes to run (default: 1).",
default=1,
)
parser.add_argument(
"-a",
"--azure",
action="store_true",
help="Use Azure identity to pass an AZURE_OPENAI_AD_TOKEN to the task environment. This is necessary when using Azure-hosted OpenAI models rather than those hosted by OpenAI.",
)
parser.add_argument(
"-e",
"--env",
@ -972,7 +972,9 @@ def run_cli(args: Sequence[str]) -> None:
)
# Get the Azure bearer token generator if a token wasn't provided and there's any evidence of using Azure
azure_token_provider = get_azure_token_provider()
azure_token_provider = None
if parsed_args.azure:
azure_token_provider = get_azure_token_provider()
# Run the scenario
if parsed_args.parallel > 1: