From 959c7e92d6595ff63e29f50e813052ff192ba9c3 Mon Sep 17 00:00:00 2001 From: ks6088ts Date: Mon, 18 Aug 2025 22:59:52 +0900 Subject: [PATCH] add Azure OpenAI config for Microsoft Entra ID --- .env.template | 1 + scripts/azure_openai_operator.py | 127 +++++++++++++++++------ template_langgraph/llms/azure_openais.py | 72 +++++++++---- 3 files changed, 152 insertions(+), 48 deletions(-) diff --git a/.env.template b/.env.template index ef98efa..d8df6fa 100644 --- a/.env.template +++ b/.env.template @@ -13,6 +13,7 @@ LANGSMITH_API_KEY="lsv2_xxx" # --------- ## Azure OpenAI Service +AZURE_OPENAI_USE_MICROSOFT_ENTRA_ID="False" AZURE_OPENAI_ENDPOINT="https://xxx.openai.azure.com/" AZURE_OPENAI_API_KEY="xxx" AZURE_OPENAI_API_VERSION="2025-04-01-preview" diff --git a/scripts/azure_openai_operator.py b/scripts/azure_openai_operator.py index a88515e..7f3f54a 100644 --- a/scripts/azure_openai_operator.py +++ b/scripts/azure_openai_operator.py @@ -1,8 +1,10 @@ import logging from base64 import b64encode +from logging import basicConfig import typer from dotenv import load_dotenv +from langchain_core.messages import HumanMessage from template_langgraph.llms.azure_openais import AzureOpenAiWrapper from template_langgraph.loggers import get_logger @@ -22,6 +24,12 @@ def load_image_to_base64(image_path: str) -> str: return b64encode(image_file.read()).decode("utf-8") +def set_verbose_logging(verbose: bool): + if verbose: + logger.setLevel(logging.DEBUG) + basicConfig(level=logging.DEBUG) + + @app.command() def chat( query: str = typer.Option( @@ -30,6 +38,12 @@ def chat( "-q", help="Query to run with the Azure OpenAI chat model", ), + stream: bool = typer.Option( + False, + "--stream", + "-s", + help="Enable streaming output", + ), verbose: bool = typer.Option( False, "--verbose", @@ -37,21 +51,36 @@ def chat( help="Enable verbose output", ), ): - # Set up logging - if verbose: - logger.setLevel(logging.DEBUG) + set_verbose_logging(verbose) logger.info("Running...") - response = AzureOpenAiWrapper().chat_model.invoke( - input=query, - ) - logger.debug( - response.model_dump_json( - indent=2, - exclude_none=True, + llm = AzureOpenAiWrapper().chat_model + + if stream: + response = "" + for chunk in llm.stream( + input=[ + HumanMessage(content=query), + ], + ): + print( + chunk.content, + end="|", + flush=True, + ) + response += str(chunk.content) + logger.info(f"Output: {response}") + else: + response = llm.invoke( + input=query, ) - ) - logger.info(f"Output: {response.content}") + logger.debug( + response.model_dump_json( + indent=2, + exclude_none=True, + ) + ) + logger.info(f"Output: {response.content}") @app.command() @@ -62,6 +91,12 @@ def reasoning( "-q", help="Query to run with the Azure OpenAI reasoning model", ), + stream: bool = typer.Option( + False, + "--stream", + "-s", + help="Enable streaming output", + ), verbose: bool = typer.Option( False, "--verbose", @@ -69,21 +104,57 @@ def reasoning( help="Enable verbose output", ), ): - # Set up logging - if verbose: - logger.setLevel(logging.DEBUG) - - logger.info("Running...") - response = AzureOpenAiWrapper().reasoning_model.invoke( - input=query, - ) - logger.debug( - response.model_dump_json( - indent=2, - exclude_none=True, + set_verbose_logging(verbose) + + llm = AzureOpenAiWrapper().reasoning_model + if stream: + response = "" + for chunk in llm.stream( + input=[ + HumanMessage(content=query), + ], + ): + print( + chunk.content, + end="|", + flush=True, + ) + response += str(chunk.content) + logger.info(f"Output: {response}") + else: + response = llm.invoke( + input=query, ) - ) - logger.info(f"Output: {response.content}") + logger.debug( + response.model_dump_json( + indent=2, + exclude_none=True, + ) + ) + logger.info(f"Output: {response.content}") + + +@app.command() +def embedding( + query: str = typer.Option( + "患者のデータから考えられる病名を診断してください。年齢： 55歳, 性別： 男性, 主訴： 激しい胸の痛み、息切れ, 検査データ： 心電図異常、トロポニン値上昇, 病歴： 高血圧、喫煙歴あり", # noqa: E501 + "--query", + "-q", + help="Query to run with the Azure OpenAI embedding model", + ), + verbose: bool = typer.Option( + False, + "--verbose", + "-v", + help="Enable verbose output", + ), +): + set_verbose_logging(verbose) + + llm = AzureOpenAiWrapper().embedding_model + + vector = llm.embed_query(text=query) + logger.info(f"Dimension: {len(vector)}, Vector: {vector[:5]}") @app.command() @@ -107,9 +178,7 @@ def image( help="Enable verbose output", ), ): - # Set up logging - if verbose: - logger.setLevel(logging.DEBUG) + set_verbose_logging(verbose) base64_image = load_image_to_base64(file_path) messages = { diff --git a/template_langgraph/llms/azure_openais.py b/template_langgraph/llms/azure_openais.py index 85dd4be..72cada7 100644 --- a/template_langgraph/llms/azure_openais.py +++ b/template_langgraph/llms/azure_openais.py @@ -1,10 +1,16 @@ from functools import lru_cache +from azure.identity import DefaultAzureCredential from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings from pydantic_settings import BaseSettings, SettingsConfigDict +from template_langgraph.loggers import get_logger + +logger = get_logger(__name__) + class Settings(BaseSettings): + azure_openai_use_microsoft_entra_id: str = "False" azure_openai_endpoint: str = "https://.openai.azure.com/" azure_openai_api_key: str = "" azure_openai_api_version: str = "2024-10-21" @@ -29,25 +35,53 @@ def __init__(self, settings: Settings = None): if settings is None: settings = get_azure_openai_settings() - self.chat_model = AzureChatOpenAI( - azure_endpoint=settings.azure_openai_endpoint, - api_key=settings.azure_openai_api_key, - api_version=settings.azure_openai_api_version, - azure_deployment=settings.azure_openai_model_chat, - streaming=True, - ) - self.embedding_model = AzureOpenAIEmbeddings( - azure_endpoint=settings.azure_openai_endpoint, - api_key=settings.azure_openai_api_key, - api_version=settings.azure_openai_api_version, - azure_deployment=settings.azure_openai_model_embedding, - ) - self.reasoning_model = AzureChatOpenAI( - azure_endpoint=settings.azure_openai_endpoint, - api_key=settings.azure_openai_api_key, - api_version=settings.azure_openai_api_version, - azure_deployment=settings.azure_openai_model_reasoning, - ) + if settings.azure_openai_use_microsoft_entra_id.lower() == "true": + logger.info("Using Microsoft Entra ID for authentication") + credential = DefaultAzureCredential() + token = credential.get_token("https://cognitiveservices.azure.com/.default").token + + self.chat_model = AzureChatOpenAI( + azure_ad_token=token, + azure_endpoint=settings.azure_openai_endpoint, + api_version=settings.azure_openai_api_version, + azure_deployment=settings.azure_openai_model_chat, + streaming=True, + ) + self.reasoning_model = AzureChatOpenAI( + azure_ad_token=token, + azure_endpoint=settings.azure_openai_endpoint, + api_version=settings.azure_openai_api_version, + azure_deployment=settings.azure_openai_model_reasoning, + streaming=True, + ) + self.embedding_model = AzureOpenAIEmbeddings( + azure_ad_token=token, + azure_endpoint=settings.azure_openai_endpoint, + api_version=settings.azure_openai_api_version, + azure_deployment=settings.azure_openai_model_embedding, + ) + else: + logger.info("Using API key for authentication") + self.chat_model = AzureChatOpenAI( + api_key=settings.azure_openai_api_key, + azure_endpoint=settings.azure_openai_endpoint, + api_version=settings.azure_openai_api_version, + azure_deployment=settings.azure_openai_model_chat, + streaming=True, + ) + self.reasoning_model = AzureChatOpenAI( + api_key=settings.azure_openai_api_key, + azure_endpoint=settings.azure_openai_endpoint, + api_version=settings.azure_openai_api_version, + azure_deployment=settings.azure_openai_model_reasoning, + streaming=True, + ) + self.embedding_model = AzureOpenAIEmbeddings( + api_key=settings.azure_openai_api_key, + azure_endpoint=settings.azure_openai_endpoint, + api_version=settings.azure_openai_api_version, + azure_deployment=settings.azure_openai_model_embedding, + ) def create_embedding(self, text: str): """Create an embedding for the given text."""