From d32a64730b1fdc04025b2731a6bf4a88d9aabae8 Mon Sep 17 00:00:00 2001
From: minggu <minggu@microsoft.com>
Date: Tue, 23 Dec 2025 16:20:33 +0800
Subject: [PATCH] Improve Azure example configuration handling

---
 examples/azure/aoai_finetune.py       | 11 ++++++++++-
 examples/azure/capital_agent.py       | 18 +++++++++++++++++-
 examples/azure/train_capital_agent.py |  6 ++++--
 3 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/examples/azure/aoai_finetune.py b/examples/azure/aoai_finetune.py
index 0296724fb..f16ad6a9c 100644
--- a/examples/azure/aoai_finetune.py
+++ b/examples/azure/aoai_finetune.py
@@ -18,7 +18,7 @@
 
 from agentlightning.adapter.messages import OpenAIMessages, TraceToMessages
 from agentlightning.algorithm import Algorithm
-from agentlightning.algorithm.utils import batch_iter_over_dataset
+from agentlightning.algorithm.apo.apo import batch_iter_over_dataset
 from agentlightning.reward import find_final_reward
 from agentlightning.types import LLM, RolloutMode, TaskInput
 
@@ -29,6 +29,11 @@
 FINETUNE_JOB_POLL_INTERVAL = 60
 
 
+from dotenv import load_dotenv
+
+# Load environment variables from .env file (override existing)
+load_dotenv(override=True)
+
 class AzureOpenAIFinetune(Algorithm):
     """Coordinate iterative fine-tuning runs for an Azure OpenAI deployment.
 
@@ -176,6 +181,10 @@ async def run(  # type: ignore
             training_data = await self.prepare_data_for_training(messages_group, reward_group, "train")
             self._log_info(f"[Stage 4] Prepared {len(training_data)} training examples after filtering.")
 
+            self.openai_client = OpenAI(
+                api_key=self.azure_openai_api_key,
+                base_url=self.azure_openai_endpoint + "/openai/v1/",
+            )
             # (5) Perform fine-tuning
             self._log_info(f"[Stage 5] Starting fine-tuning for model {training_model_name}...")
             training_model_name = self.finetune(training_data, training_model_name, i_iteration)
diff --git a/examples/azure/capital_agent.py b/examples/azure/capital_agent.py
index c9a496448..5c215bee1 100644
--- a/examples/azure/capital_agent.py
+++ b/examples/azure/capital_agent.py
@@ -28,6 +28,13 @@
 
 from agentlightning import LLM, AgentOpsTracer, InMemoryLightningStore, LitAgentRunner, rollout
 
+from dotenv import load_dotenv
+
+# Load environment variables from .env file (override existing)
+load_dotenv(override=True)
+
+
+
 CAPITALS = {
     "japan": "Tokyo",
     "france": "Paris",
@@ -84,7 +91,16 @@ def capital_agent(task: CapitalTask, llm: LLM) -> float:
     prompt = task["input"]
     expected = task["output"]
 
-    openai_client = openai.OpenAI(base_url=llm.endpoint, api_key=os.getenv("AZURE_OPENAI_API_KEY", ""))
+    api_version = os.getenv("AZURE_OPENAI_API_VERSION", "")
+    if not api_version:
+        raise ValueError("Set AZURE_OPENAI_API_VERSION to match your Azure OpenAI resource.")
+
+    # Use AzureOpenAI client for Azure endpoints
+    openai_client = openai.AzureOpenAI(
+        azure_endpoint=llm.endpoint,
+        api_key=os.getenv("AZURE_OPENAI_API_KEY", ""),
+        api_version=api_version,
+    )
 
     messages: List[ChatCompletionMessageParam] = [
         {"role": "system", "content": SYSTEM},
diff --git a/examples/azure/train_capital_agent.py b/examples/azure/train_capital_agent.py
index 39b07f7c8..058530a5c 100644
--- a/examples/azure/train_capital_agent.py
+++ b/examples/azure/train_capital_agent.py
@@ -7,7 +7,9 @@
 from capital_agent import capital_agent
 from rich.console import Console
 
-from agentlightning import TraceToMessages, Trainer, setup_logging
+from agentlightning.adapter.messages import TraceToMessages
+from agentlightning.trainer.trainer import Trainer
+from agentlightning.logging import configure_logger
 
 console = Console()
 
@@ -20,7 +22,7 @@ def parse_args() -> argparse.Namespace:
 
 
 def main():
-    setup_logging()
+    configure_logger()
     args = parse_args()
     finetune_algo = AzureOpenAIFinetune(
         base_deployment_name="gpt-4.1-mini",