Python: Fix Eval samples (#4033)

* fix red team sample

* Updated self-reflection

* fix for workflow eval sample

* fix test

python/samples/05-end-to-end/evaluation/red_teaming/red_team_agent_sample.py

@@ -1,3 +1,13 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "azure-ai-evaluation",
+#     "pyrit==0.9.0"
+# ]
+# ///
+# Run with any PEP 723 compatible runner, e.g.:
+#   uv run samples/05-end-to-end/evaluation/red_teaming/red_team_agent_sample.py
+
 # Copyright (c) Microsoft. All rights reserved.
 # type: ignore
 import asyncio
@@ -5,6 +15,7 @@ import json
 import os
 from typing import Any
 
+from agent_framework import Message
 from agent_framework.azure import AzureOpenAIChatClient
 from azure.ai.evaluation.red_team import AttackStrategy, RedTeam, RiskCategory
 from azure.identity import AzureCliCredential
@@ -20,10 +31,10 @@ the safety and resilience of an Agent Framework agent against adversarial attack
 Prerequisites:
     - Azure AI project (hub and project created)
     - Azure CLI authentication (run `az login`)
-    - Environment variables set in .env file or environment
+    - Environment variables set in environment
 
 Installation:
-    pip install agent-framework azure-ai-evaluation pyrit duckdb azure-identity
+    pip install agent-framework-core azure-ai-evaluation pyrit==0.9.0 duckdb
 
 Reference:
     Azure AI Red Teaming: https://github.com/Azure-Samples/azureai-samples/blob/main/scenarios/evaluate/AI_RedTeaming/AI_RedTeaming.ipynb
@@ -60,19 +71,30 @@ Your boundaries:
     )
 
     # Create the callback
-    async def agent_callback(query: str) -> dict[str, list[Any]]:
+    async def agent_callback(
+        messages: list,
+        stream: bool | None = False,  # noqa: ARG001
+        session_state: str | None = None,  # noqa: ARG001
+        context: dict[str, Any] | None = None,  # noqa: ARG001
+    ) -> dict[str, list[dict[str, str]]]:
         """Async callback function that interfaces between RedTeam and the agent.
 
         Args:
-            query: The adversarial prompt from RedTeam
+            messages: The adversarial prompts from RedTeam
         """
+        messages_list = [Message(role=message.role, text=message.content) for message in messages]
         try:
-            response = await agent.run(query)
-            return {"messages": [{"content": response.text, "role": "assistant"}]}
-
+            response = agent.run(messages=messages_list, stream=stream)
+            result = await response.get_final_response() if stream else await response
+            # Format the response to follow the expected chat protocol format
+            formatted_response = {"content": result.text, "role": "assistant"}
         except Exception as e:
-            print(f"Error during agent run: {e}")
-            return {"messages": [f"I encountered an error and couldn't process your request: {e!s}"]}
+            print(f"Error calling Azure OpenAI: {e!s}")
+            formatted_response = {
+                "content": f"I encountered an error and couldn't process your request: {e}",
+                "role": "assistant",
+            }
+        return {"messages": [formatted_response]}
 
     # Create RedTeam instance
     red_team = RedTeam(

python/samples/05-end-to-end/evaluation/self_reflection/README.md

@@ -7,23 +7,22 @@ This sample demonstrates the self-reflection pattern using Agent Framework and A
 **What it demonstrates:**
 - Iterative self-reflection loop that automatically improves responses based on groundedness evaluation
 - Batch processing of prompts from JSONL files with progress tracking
-- Using `AzureOpenAIChatClient` with Azure CLI authentication
+- Using `AzureOpenAIResponsesClient` with a Project Endpoint and Azure CLI authentication
 - Comprehensive summary statistics and detailed result tracking
 
 ## Prerequisites
 
 ### Azure Resources
-- **Azure OpenAI**: Deploy models (default: gpt-4.1 for both agent and judge)
+- **Azure OpenAI Responses in Foundry**: Deploy models (default: gpt-5.2 for both agent and judge)
 - **Azure CLI**: Run `az login` to authenticate
 
 ### Python Environment
 ```bash
-pip install agent-framework-core azure-ai-projects pandas --pre
+pip install agent-framework-core pandas --pre
 ```
 
 ### Environment Variables
 ```bash
-# .env file
 AZURE_AI_PROJECT_ENDPOINT=https://<your-ai-resource>.services.ai.azure.com/api/projects/<your-ai-project>/
 ```
 
@@ -67,6 +66,12 @@ The agent iteratively improves responses:
   ✓ Completed with score: 5/5 (best at iteration 2/3)
 ```
 
+In the Foundry UI, under `Build`/`Evaluations` you can view detailed results for each prompt, including:
+- Context
+- Query
+- Response
+- Groundedness scores and reasoning for each interation of each prompt
+
 ## Related Resources
 
 - [Reflexion Paper](https://arxiv.org/abs/2303.11366)

python/samples/05-end-to-end/evaluation/self_reflection/self_reflection.py

@@ -2,6 +2,7 @@
 # requires-python = ">=3.10"
 # dependencies = [
 #     "pandas",
+#     "pyarrow",
 # ]
 # ///
 # Run with any PEP 723 compatible runner, e.g.:
@@ -13,12 +14,13 @@ import argparse
 import asyncio
 import os
 import time
+from pathlib import Path
 from typing import Any
 
 import openai
 import pandas as pd
 from agent_framework import Agent, Message
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.azure import AzureOpenAIResponsesClient
 from azure.ai.projects import AIProjectClient
 from azure.identity import AzureCliCredential
 from dotenv import load_dotenv
@@ -50,11 +52,38 @@ Usage as CLI with extra options:
                               --output resources/results.jsonl \\
                               --max-reflections 3 \\
                               -n 10  # Optional: process only first 10 prompts
+
+=============== Example output ===============
+
+============================================================
+SUMMARY
+============================================================
+Total prompts processed: 31
+  ✓ Successful: 30
+  ✗ Failed: 1
+
+Groundedness Scores:
+  Average best score: 4.77/5
+  Perfect scores (5/5): 25/30 (83.3%)
+
+Improvement Analysis:
+  Average first score: 4.50/5
+  Average final score: 4.70/5
+  Average improvement: +0.20
+  Responses that improved: 4/30 (13.3%)
+
+Iteration Statistics:
+  Average best iteration: 1.17
+  Best on first try: 25/30 (83.3%)
+============================================================
+
+✓ Processing complete!
+
 """
 
 
-DEFAULT_AGENT_MODEL = "gpt-4.1"
-DEFAULT_JUDGE_MODEL = "gpt-4.1"
+DEFAULT_AGENT_MODEL = "gpt-5.2"
+DEFAULT_JUDGE_MODEL = "gpt-5.2"
 
 
 def create_openai_client():
@@ -64,6 +93,13 @@ def create_openai_client():
     return project_client.get_openai_client()
 
 
+def create_async_project_client():
+    from azure.ai.projects.aio import AIProjectClient as AsyncAIProjectClient
+    from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential
+
+    return AsyncAIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=AsyncAzureCliCredential())
+
+
 def create_eval(client: openai.OpenAI, judge_model: str) -> openai.types.EvalCreateResponse:
     print("Creating Eval")
     data_source_config = DataSourceConfigCustom({
@@ -257,6 +293,7 @@ async def execute_query_with_self_reflection(
 
 
 async def run_self_reflection_batch(
+    project_client: AIProjectClient,
     input_file: str,
     output_file: str,
     agent_model: str = DEFAULT_AGENT_MODEL,
@@ -284,16 +321,15 @@ async def run_self_reflection_batch(
         load_dotenv(override=True)
 
     # Create agent, it loads environment variables AZURE_OPENAI_API_KEY and AZURE_OPENAI_ENDPOINT automatically
-    agent = AzureOpenAIChatClient(
-        credential=AzureCliCredential(),
+    responses_client = AzureOpenAIResponsesClient(
+        project_client=project_client,
         deployment_name=agent_model,
-    ).as_agent(
-        instructions="You are a helpful agent.",
     )
 
     # Load input data
-    print(f"Loading prompts from: {input_file}")
-    df = pd.read_json(input_file, lines=True)
+    input_path = (Path(__file__).parent / input_file).resolve()
+    print(f"Loading prompts from: {input_path}")
+    df = pd.read_json(path_or_buf=input_path, lines=True, engine="pyarrow")
     print(f"Loaded {len(df)} prompts")
 
     # Apply limit if specified
@@ -332,7 +368,7 @@ async def run_self_reflection_batch(
         try:
             result = await execute_query_with_self_reflection(
                 client=client,
-                agent=agent,
+                agent=responses_client.as_agent(instructions=row["system_instruction"]),
                 eval_object=eval_object,
                 full_user_query=row["full_prompt"],
                 context=row["context_document"],
@@ -386,8 +422,9 @@ async def run_self_reflection_batch(
     # Create DataFrame and save
     results_df = pd.DataFrame(results)
 
-    print(f"\nSaving results to: {output_file}")
-    results_df.to_json(output_file, orient="records", lines=True)
+    output_path = (Path(__file__).parent / output_file).resolve()
+    print(f"\nSaving results to: {output_path}")
+    results_df.to_json(output_path, orient="records", lines=True)
 
     # Generate detailed summary
     successful_runs = results_df[results_df["error"].isna()]
@@ -482,6 +519,7 @@ async def main():
     # Run the batch processing
     try:
         await run_self_reflection_batch(
+            project_client=create_async_project_client(),
             input_file=args.input,
             output_file=args.output,
             agent_model=args.agent_model,
@@ -499,4 +537,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    exit(asyncio.run(main()))
+    asyncio.run(main())

python/samples/05-end-to-end/workflow_evaluation/.env.example

@@ -1,2 +1,3 @@
 AZURE_AI_PROJECT_ENDPOINT="<your-project-endpoint>"
-AZURE_AI_MODEL_DEPLOYMENT_NAME="<your-model-deployment>"
+AZURE_AI_MODEL_DEPLOYMENT_NAME_WORKFLOW="<your-model-deployment>"
+AZURE_AI_MODEL_DEPLOYMENT_NAME_EVAL="<your-model-deployment>"

python/samples/05-end-to-end/workflow_evaluation/create_workflow.py

@@ -1,5 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
-
+# type: ignore
 """
 Multi-Agent Travel Planning Workflow Evaluation with Multiple Response Tracking
 
@@ -52,10 +52,11 @@ from agent_framework import (
     Message,
     WorkflowBuilder,
     WorkflowContext,
+    WorkflowEvent,
     executor,
     handler,
 )
-from agent_framework.azure import AzureAIClient
+from agent_framework.azure import AzureOpenAIResponsesClient
 from azure.ai.projects.aio import AIProjectClient
 from azure.identity.aio import DefaultAzureCredential
 from dotenv import load_dotenv
@@ -73,8 +74,8 @@ async def start_executor(input: str, ctx: WorkflowContext[list[Message]]) -> Non
 class ResearchLead(Executor):
     """Aggregates and summarizes travel planning findings from all specialized agents."""
 
-    def __init__(self, client: AzureAIClient, id: str = "travel-planning-coordinator"):
-        # store=True to preserve conversation history for evaluation
+    def __init__(self, client: AzureOpenAIResponsesClient, id: str = "travel-planning-coordinator"):
+        # Use default_options to persist conversation history for evaluation.
         self.agent = client.as_agent(
             id="travel-planning-coordinator",
             instructions=(
@@ -86,7 +87,6 @@ class ResearchLead(Executor):
                 "Clearly indicate which information came from which agent. Do not use tools."
             ),
             name="travel-planning-coordinator",
-            store=True,
         )
         super().__init__(id=id)
 
@@ -142,12 +142,15 @@ class ResearchLead(Executor):
         return agent_findings
 
 
-async def run_workflow_with_response_tracking(query: str, client: AzureAIClient | None = None) -> dict:
+async def run_workflow_with_response_tracking(
+    query: str, client: AzureOpenAIResponsesClient | None = None, deployment_name: str | None = None
+) -> dict:
     """Run multi-agent workflow and track conversation IDs, response IDs, and interaction sequence.
 
     Args:
         query: The user query to process through the multi-agent workflow
-        client: Optional AzureAIClient instance
+        client: Optional AzureOpenAIResponsesClient instance
+        deployment_name: Optional model deployment name for the workflow agents
 
     Returns:
         Dictionary containing interaction sequence, conversation/response IDs, and conversation analysis
@@ -155,17 +158,13 @@ async def run_workflow_with_response_tracking(query: str, client: AzureAIClient
     if client is None:
         try:
             async with DefaultAzureCredential() as credential:
-                # Create AIProjectClient with the correct API version for V2 prompt agents
                 project_client = AIProjectClient(
                     endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
                     credential=credential,
-                    api_version="2025-11-15-preview",
                 )
 
-                async with (
-                    project_client,
-                    AzureAIClient(project_client=project_client, credential=credential) as client,
-                ):
+                async with project_client:
+                    client = AzureOpenAIResponsesClient(project_client=project_client, deployment_name=deployment_name)
                     return await _run_workflow_with_client(query, client)
         except Exception as e:
             print(f"Error during workflow execution: {e}")
@@ -174,21 +173,29 @@ async def run_workflow_with_response_tracking(query: str, client: AzureAIClient
         return await _run_workflow_with_client(query, client)
 
 
-async def _run_workflow_with_client(query: str, client: AzureAIClient) -> dict:
+async def _run_workflow_with_client(query: str, client: AzureOpenAIResponsesClient) -> dict:
     """Execute workflow with given client and track all interactions."""
 
     # Initialize tracking variables - use lists to track multiple responses per agent
-    conversation_ids = defaultdict(list)
-    response_ids = defaultdict(list)
-    workflow_output = None
+    conversation_ids: dict[str, list[str]] = defaultdict(list)
+    response_ids: dict[str, list[str]] = defaultdict(list)
 
-    # Create workflow components and keep agent references
-    # Pass project_client and credential to create separate client instances per agent
-    workflow, agent_map = await _create_workflow(client.project_client, client.credential)
+    # Create workflow components using a single shared client
+    workflow, agent_map = await _create_workflow(client)
 
-    # Process workflow events
-    events = workflow.run(query, stream=True)
-    workflow_output = await _process_workflow_events(events, conversation_ids, response_ids)
+    def track_ids(event: WorkflowEvent) -> WorkflowEvent:
+        """Transform hook that tracks response/conversation IDs from AgentResponseUpdate events."""
+        if event.type == "output" and isinstance(event.data, AgentResponseUpdate):
+            _track_agent_ids(event, event.executor_id, response_ids, conversation_ids)
+        return event
+
+    # Process workflow events using a transform hook for ID tracking
+    stream = workflow.run(query, stream=True).with_transform_hook(track_ids)
+    result = await stream.get_final_response()
+
+    workflow_output = result.get_outputs()[-1] if result.get_outputs() else None
+    if workflow_output:
+        print(f"\nWorkflow Output: {workflow_output}\n")
 
     return {
         "conversation_ids": dict(conversation_ids),
@@ -198,115 +205,80 @@ async def _run_workflow_with_client(query: str, client: AzureAIClient) -> dict:
     }
 
 
-async def _create_workflow(project_client, credential):
+async def _create_workflow(client: AzureOpenAIResponsesClient):
     """Create the multi-agent travel planning workflow with specialized agents.
 
-    IMPORTANT: Each agent needs its own client instance because the V2 client stores
-    agent_name and agent_version as instance variables, causing all agents to share
-    the same agent identity if they share a client.
+    Uses a single shared AzureOpenAIResponsesClient for all agents.
     """
 
-    # Create separate client for Final Coordinator
-    final_coordinator_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="final-coordinator"
-    )
-    final_coordinator = ResearchLead(client=final_coordinator_client, id="final-coordinator")
+    final_coordinator = ResearchLead(client=client, id="final-coordinator")
 
     # Agent 1: Travel Request Handler (initial coordinator)
-    # Create separate client with unique agent_name
-    travel_request_handler_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="travel-request-handler"
-    )
-    travel_request_handler = travel_request_handler_client.as_agent(
+    travel_request_handler = client.as_agent(
         id="travel-request-handler",
         instructions=(
             "You receive user travel queries and relay them to specialized agents. Extract key information: destination, dates, budget, and preferences. Pass this information forward clearly to the next agents."
         ),
         name="travel-request-handler",
-        store=True,
     )
 
     # Agent 2: Hotel Search Executor
-    hotel_search_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="hotel-search-agent"
-    )
-    hotel_search_agent = hotel_search_client.as_agent(
+    hotel_search_agent = client.as_agent(
         id="hotel-search-agent",
         instructions=(
             "You are a hotel search specialist. Your task is ONLY to search for and provide hotel information. Use search_hotels to find options, get_hotel_details for specifics, and check_availability to verify rooms. Output format: List hotel names, prices per night, total cost for the stay, locations, ratings, amenities, and addresses. IMPORTANT: Only provide hotel information without additional commentary."
         ),
         name="hotel-search-agent",
         tools=[search_hotels, get_hotel_details, check_hotel_availability],
-        store=True,
     )
 
     # Agent 3: Flight Search Executor
-    flight_search_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="flight-search-agent"
-    )
-    flight_search_agent = flight_search_client.as_agent(
+    flight_search_agent = client.as_agent(
         id="flight-search-agent",
         instructions=(
             "You are a flight search specialist. Your task is ONLY to search for and provide flight information. Use search_flights to find options, get_flight_details for specifics, and check_availability for seats. Output format: List flight numbers, airlines, departure/arrival times, prices, durations, and cabin class. IMPORTANT: Only provide flight information without additional commentary."
         ),
         name="flight-search-agent",
         tools=[search_flights, get_flight_details, check_flight_availability],
-        store=True,
     )
 
     # Agent 4: Activity Search Executor
-    activity_search_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="activity-search-agent"
-    )
-    activity_search_agent = activity_search_client.as_agent(
+    activity_search_agent = client.as_agent(
         id="activity-search-agent",
         instructions=(
             "You are an activities specialist. Your task is ONLY to search for and provide activity information. Use search_activities to find options for activities. Output format: List activity names, descriptions, prices, durations, ratings, and categories. IMPORTANT: Only provide activity information without additional commentary."
         ),
         name="activity-search-agent",
         tools=[search_activities],
-        store=True,
     )
 
     # Agent 5: Booking Confirmation Executor
-    booking_confirmation_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="booking-confirmation-agent"
-    )
-    booking_confirmation_agent = booking_confirmation_client.as_agent(
+    booking_confirmation_agent = client.as_agent(
         id="booking-confirmation-agent",
         instructions=(
             "You confirm bookings. Use check_hotel_availability and check_flight_availability to verify slots, then confirm_booking to finalize. Provide ONLY: confirmation numbers, booking references, and confirmation status."
         ),
         name="booking-confirmation-agent",
         tools=[confirm_booking, check_hotel_availability, check_flight_availability],
-        store=True,
     )
 
     # Agent 6: Booking Payment Executor
-    booking_payment_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="booking-payment-agent"
-    )
-    booking_payment_agent = booking_payment_client.as_agent(
+    booking_payment_agent = client.as_agent(
         id="booking-payment-agent",
         instructions=(
             "You process payments. Use validate_payment_method to verify payment, then process_payment to complete transactions. Provide ONLY: payment confirmation status, transaction IDs, and payment amounts."
         ),
         name="booking-payment-agent",
         tools=[process_payment, validate_payment_method],
-        store=True,
     )
 
     # Agent 7: Booking Information Aggregation Executor
-    booking_info_client = AzureAIClient(
-        project_client=project_client, credential=credential, agent_name="booking-info-aggregation-agent"
-    )
-    booking_info_aggregation_agent = booking_info_client.as_agent(
+    booking_info_aggregation_agent = client.as_agent(
         id="booking-info-aggregation-agent",
         instructions=(
             "You aggregate hotel and flight search results. Receive options from search agents and organize them. Provide: top 2-3 hotel options with prices and top 2-3 flight options with prices in a structured format."
         ),
         name="booking-info-aggregation-agent",
-        store=True,
     )
 
     # Build workflow with logical booking flow:
@@ -347,63 +319,31 @@ async def _create_workflow(project_client, credential):
     return workflow, agent_map
 
 
-async def _process_workflow_events(events, conversation_ids, response_ids):
-    """Process workflow events and track interactions."""
-    workflow_output = None
-
-    async for event in events:
-        if event.type == "output":
-            workflow_output = event.data
-            # Handle Unicode characters that may not be displayable in Windows console
-            try:
-                print(f"\nWorkflow Output: {event.data}\n")
-            except UnicodeEncodeError:
-                output_str = str(event.data).encode("ascii", "replace").decode("ascii")
-                print(f"\nWorkflow Output: {output_str}\n")
-
-        elif event.type == "output" and isinstance(event.data, AgentResponseUpdate):
-            _track_agent_ids(event, event.executor_id, response_ids, conversation_ids)
-
-    return workflow_output
-
-
 def _track_agent_ids(event, agent, response_ids, conversation_ids):
     """Track agent response and conversation IDs - supporting multiple responses per agent."""
+    update = event.data
+
+    # response_id is directly on AgentResponseUpdate
+    if update.response_id and update.response_id not in response_ids[agent]:
+        response_ids[agent].append(update.response_id)
+
+    # conversation_id is on the underlying ChatResponseUpdate (raw_representation)
+    raw = update.raw_representation
     if (
-        isinstance(event.data, AgentResponseUpdate)
-        and hasattr(event.data, "raw_representation")
-        and event.data.raw_representation
+        raw
+        and hasattr(raw, "conversation_id")
+        and raw.conversation_id
+        and raw.conversation_id not in conversation_ids[agent]
     ):
-        # Check for conversation_id and response_id from raw_representation
-        # V2 API stores conversation_id directly on raw_representation (ChatResponseUpdate)
-        raw = event.data.raw_representation
-
-        # Try conversation_id directly on raw representation
-        if (
-            hasattr(raw, "conversation_id")
-            and raw.conversation_id  # type: ignore[union-attr]
-            and raw.conversation_id not in conversation_ids[agent]  # type: ignore[union-attr]
-        ):
-            # Only add if not already in the list
-            conversation_ids[agent].append(raw.conversation_id)  # type: ignore[union-attr]
-
-        # Extract response_id from the OpenAI event (available from first event)
-        if hasattr(raw, "raw_representation") and raw.raw_representation:  # type: ignore[union-attr]
-            openai_event = raw.raw_representation  # type: ignore[union-attr]
-
-            # Check if event has response object with id
-            if (
-                hasattr(openai_event, "response")
-                and hasattr(openai_event.response, "id")
-                and openai_event.response.id not in response_ids[agent]
-            ):
-                # Only add if not already in the list
-                response_ids[agent].append(openai_event.response.id)
+        conversation_ids[agent].append(raw.conversation_id)
 
 
-async def create_and_run_workflow():
+async def create_and_run_workflow(deployment_name: str | None = None):
     """Run the workflow evaluation and display results.
 
+    Args:
+        deployment_name: Optional model deployment name for the workflow agents
+
     Returns:
         Dictionary containing agents data with conversation IDs, response IDs, and query information
     """
@@ -416,7 +356,7 @@ async def create_and_run_workflow():
     query = example_queries[0]
     print(f"Query: {query}\n")
 
-    result = await run_workflow_with_response_tracking(query)
+    result = await run_workflow_with_response_tracking(query, deployment_name=deployment_name)
 
     # Create output data structure
     output_data = {"agents": {}, "query": result["query"], "output": result.get("output", "")}

python/samples/05-end-to-end/workflow_evaluation/run_evaluation.py

@@ -1,24 +1,43 @@
 # Copyright (c) Microsoft. All rights reserved.
+# type: ignore
 
-"""
-Script to run multi-agent travel planning workflow and evaluate agent responses.
-
-This script:
-1. Executes the multi-agent workflow
-2. Displays response data summary
-3. Creates and runs evaluation with multiple evaluators
-4. Monitors evaluation progress and displays results
-"""
+from __future__ import annotations
 
 import asyncio
 import os
 import time
+from typing import TYPE_CHECKING, Any
 
 from azure.ai.projects import AIProjectClient
 from azure.identity import DefaultAzureCredential
 from create_workflow import create_and_run_workflow
 from dotenv import load_dotenv
 
+if TYPE_CHECKING:
+    from openai import OpenAI
+    from openai.types import EvalCreateResponse
+    from openai.types.evals import RunCreateResponse
+
+"""
+Script to run multi-agent travel planning workflow and evaluate agent responses.
+
+This script:
+1. Runs the multi-agent travel planning workflow
+2. Displays a summary of tracked agent responses
+3. Fetches and previews final agent responses
+4. Creates an evaluation with multiple evaluators
+5. Runs the evaluation on selected agent responses
+6. Monitors evaluation progress and displays results
+"""
+
+
+def create_openai_client() -> OpenAI:
+    project_client = AIProjectClient(
+        endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
+        credential=DefaultAzureCredential(),
+    )
+    return project_client.get_openai_client()
+
 
 def print_section(title: str):
     """Print a formatted section header."""
@@ -27,26 +46,26 @@ def print_section(title: str):
     print(f"{'=' * 80}")
 
 
-async def run_workflow():
+async def run_workflow(deployment_name: str | None = None) -> dict[str, Any]:
     """Execute the multi-agent travel planning workflow.
 
+    Args:
+        deployment_name: Optional model deployment name for the workflow agents
+
     Returns:
         Dictionary containing workflow data with agent response IDs
     """
-    print_section("Step 1: Running Workflow")
     print("Executing multi-agent travel planning workflow...")
     print("This may take a few minutes...")
 
-    workflow_data = await create_and_run_workflow()
+    workflow_data = await create_and_run_workflow(deployment_name=deployment_name)
 
     print("Workflow execution completed")
     return workflow_data
 
 
-def display_response_summary(workflow_data: dict):
+def display_response_summary(workflow_data: dict) -> None:
     """Display summary of response data."""
-    print_section("Step 2: Response Data Summary")
-
     print(f"Query: {workflow_data['query']}")
     print(f"\nAgents tracked: {len(workflow_data['agents'])}")
 
@@ -55,10 +74,8 @@ def display_response_summary(workflow_data: dict):
         print(f"  {agent_name}: {response_count} response(s)")
 
 
-def fetch_agent_responses(openai_client, workflow_data: dict, agent_names: list):
+def fetch_agent_responses(openai_client: OpenAI, workflow_data: dict[str, Any], agent_names: list[str]) -> None:
     """Fetch and display final responses from specified agents."""
-    print_section("Step 3: Fetching Agent Responses")
-
     for agent_name in agent_names:
         if agent_name not in workflow_data["agents"]:
             continue
@@ -80,10 +97,9 @@ def fetch_agent_responses(openai_client, workflow_data: dict, agent_names: list)
             print(f"  Error: {e}")
 
 
-def create_evaluation(openai_client, model_deployment: str):
+def create_evaluation(openai_client: OpenAI, deployment_name: str | None = "gpt-5.2") -> EvalCreateResponse:
     """Create evaluation with multiple evaluators."""
-    print_section("Step 4: Creating Evaluation")
-
+    deployment_name = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME", deployment_name)
     data_source_config = {"type": "azure_ai_source", "scenario": "responses"}
 
     testing_criteria = [
@@ -91,25 +107,25 @@ def create_evaluation(openai_client, model_deployment: str):
             "type": "azure_ai_evaluator",
             "name": "relevance",
             "evaluator_name": "builtin.relevance",
-            "initialization_parameters": {"deployment_name": model_deployment}
+            "initialization_parameters": {"deployment_name": deployment_name},
         },
         {
             "type": "azure_ai_evaluator",
             "name": "groundedness",
             "evaluator_name": "builtin.groundedness",
-            "initialization_parameters": {"deployment_name": model_deployment}
+            "initialization_parameters": {"deployment_name": deployment_name},
         },
         {
             "type": "azure_ai_evaluator",
             "name": "tool_call_accuracy",
             "evaluator_name": "builtin.tool_call_accuracy",
-            "initialization_parameters": {"deployment_name": model_deployment}
+            "initialization_parameters": {"deployment_name": deployment_name},
         },
         {
             "type": "azure_ai_evaluator",
             "name": "tool_output_utilization",
             "evaluator_name": "builtin.tool_output_utilization",
-            "initialization_parameters": {"deployment_name": model_deployment}
+            "initialization_parameters": {"deployment_name": deployment_name},
         },
     ]
 
@@ -126,10 +142,10 @@ def create_evaluation(openai_client, model_deployment: str):
     return eval_object
 
 
-def run_evaluation(openai_client, eval_object, workflow_data: dict, agent_names: list):
+def run_evaluation(
+    openai_client: OpenAI, eval_object: EvalCreateResponse, workflow_data: dict[str, Any], agent_names: list[str]
+) -> RunCreateResponse:
     """Run evaluation on selected agent responses."""
-    print_section("Step 5: Running Evaluation")
-
     selected_response_ids = []
     for agent_name in agent_names:
         if agent_name in workflow_data["agents"]:
@@ -162,10 +178,8 @@ def run_evaluation(openai_client, eval_object, workflow_data: dict, agent_names:
     return eval_run
 
 
-def monitor_evaluation(openai_client, eval_object, eval_run):
+def monitor_evaluation(openai_client: OpenAI, eval_object: EvalCreateResponse, eval_run: RunCreateResponse):
     """Monitor evaluation progress and display results."""
-    print_section("Step 6: Monitoring Evaluation")
-
     print("Waiting for evaluation to complete...")
 
     while eval_run.status not in ["completed", "failed"]:
@@ -187,29 +201,41 @@ def monitor_evaluation(openai_client, eval_object, eval_run):
 async def main():
     """Main execution flow."""
     load_dotenv()
+    openai_client = create_openai_client()
 
-    print("Travel Planning Workflow Evaluation")
+    # Model configuration
+    workflow_agent_model = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME_WORKFLOW", "gpt-4.1-nano")
+    eval_model = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME_EVAL", "gpt-5.2")
 
-    workflow_data = await run_workflow()
+    # Focus on these agents, uncomment other ones you want to have evals run on
+    agents_to_evaluate = [
+        "hotel-search-agent",
+        "flight-search-agent",
+        "activity-search-agent",
+        # "booking-payment-agent",
+        # "booking-info-aggregation-agent",
+        # "travel-request-handler",
+        # "booking-confirmation-agent",
+    ]
 
+    print_section("Travel Planning Workflow Evaluation")
+
+    print_section("Step 1: Running Workflow")
+    workflow_data = await run_workflow(deployment_name=workflow_agent_model)
+
+    print_section("Step 2: Response Data Summary")
     display_response_summary(workflow_data)
 
-    project_client = AIProjectClient(
-        endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"],
-        credential=DefaultAzureCredential(),
-        api_version="2025-11-15-preview"
-    )
-    openai_client = project_client.get_openai_client()
-
-    agents_to_evaluate = ["hotel-search-agent", "flight-search-agent", "activity-search-agent"]
-
+    print_section("Step 3: Fetching Agent Responses")
     fetch_agent_responses(openai_client, workflow_data, agents_to_evaluate)
 
-    model_deployment = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME", "gpt-4o-mini")
-    eval_object = create_evaluation(openai_client, model_deployment)
+    print_section("Step 4: Creating Evaluation")
+    eval_object = create_evaluation(openai_client, deployment_name=eval_model)
 
+    print_section("Step 5: Running Evaluation")
     eval_run = run_evaluation(openai_client, eval_object, workflow_data, agents_to_evaluate)
 
+    print_section("Step 6: Monitoring Evaluation")
     monitor_evaluation(openai_client, eval_object, eval_run)
 
     print_section("Complete")