microsoft
How to use agentlighning with azureopenai
How to use agentlighning with azureopenai
agentlightning_azure_example.py
import os from typing import Any, Dict
import agentlightning as agl from openai import AsyncAzureOpenAI from config import Config as cfg
============================================================
1. Load Config
============================================================
cfg = Config()
if not (cfg.OPENAI_API_KEY and cfg.OPENAI_API_BASE and cfg.deployment_name): raise RuntimeError( "Missing config. Set OPENAI_API_KEY, OPENAI_API_BASE, deployment_name" )
============================================================
2. Azure Client (correct for GPT-4.1 Azure)
============================================================
client = AsyncAzureOpenAI( api_key=cfg.OPENAI_API_KEY, azure_endpoint=cfg.OPENAI_API_BASE, api_version=cfg.OPENAI_API_VERSION, )
============================================================
3. Rollout Function (Agent Behavior)
— updated to match Agent-Lightning official design
============================================================
@agl.rollout async def room_selector(task: Dict[str, Any], prompt_template: agl.PromptTemplate) -> float: """ Runs one episode of the agent. """
# Build prompt
prompt = prompt_template.format(**task)
print("************************************************")
print("AL calling with model:", cfg.deployment_name)
messages = [
{"role": "user", "content": prompt}
]
# AzureOpenAI call (must use deployment_name)
resp = await client.chat.completions.create(
model=cfg.deployment_name, # GPT-4.1 or your deployment id
messages=messages,
max_tokens=50,
temperature=0.0,
)
print(resp)
output = resp.choices[0].message.content.strip()
# Reward = 1 if correct, else 0
expected = task.get("expected_choice")
reward = 1.0 if expected and output == expected else 0.0
return reward
============================================================
4. Initial Prompt Template (APO optimizes this)
============================================================
initial_template = agl.PromptTemplate( template=( "You are a meeting-room selection assistant.\n" "Given: {people} people, time {time}, whiteboard required: {whiteboard}.\n" "Respond with ONLY the best room ID (e.g., A103)." ), engine="f-string", ) async def llm_proxy(prompt: str): resp = await client.chat.completions.create( model=cfg.deployment_name, # your Azure deployment messages=[{"role": "user", "content": prompt}], max_tokens=200, temperature=0.0, ) return resp.choices[0].message.content
============================================================
5. APO + Trainer (updated per Microsoft example)
============================================================
Create APO optimizer (uses Azure client)
algo = agl.APO(client )
algo.async_openai_client = client # ensure internal AE calls use same client algo.model = cfg.deployment_name
trainer = agl.Trainer( algorithm=algo, n_runners=1, initial_resources={"prompt_template": initial_template}, adapter=agl.TraceToMessages(),
)
============================================================
6. Training Data
============================================================
train_dataset = [ {"people": 4, "time": "10:00", "whiteboard": True, "expected_choice": "A103"}, {"people": 2, "time": "14:00", "whiteboard": False, "expected_choice": "B201"}, {"people": 6, "time": "09:00", "whiteboard": True, "expected_choice": "A103"}, {"people": 1, "time": "16:00", "whiteboard": False, "expected_choice": "C310"}, ]
val_dataset = [ {"people": 3, "time": "11:00", "whiteboard": True, "expected_choice": "A101"}, ]
============================================================
7. Train the Agent (same as Microsoft tutorial)
============================================================
trainer.fit( agent=room_selector, train_dataset=train_dataset, val_dataset=val_dataset, )
============================================================
8. Test the fine-tuned prompt
============================================================
best_prompt = trainer.resources["prompt_template"]
test_task = {"people": 5, "time": "11:00", "whiteboard": True, "expected_choice": "A103"}
final_reward = room_selector(test_task, best_prompt) print("Final Test Reward:", final_reward)
Direct raw LLM call
resp = client.chat.completions.create(
model=cfg.deployment_name,
messages=[{"role": "user", "content": best_prompt.format(**test_task)}],
max_tokens=50,
temperature=0.0,
)
print("Model Output:", resp.choices[0].message.content.strip())
this is the code I am using - facing issues
I think modern Azure OpenAI can be directly used with standard OpenAI client. No need for AzureOpenAI any more.
