foundry_multi_agent_tracing_patterns.py

"""
Foundry Multi-Agent Tracing Patterns
======================================

Demonstrates the **multi-agent observability** semantic conventions from
the Microsoft docs (OTel GenAI agent-spans spec) using the ``azure-ai-agents``
SDK.

What this adds over the single-agent demos:
  - Multiple collaborating agents (Router → Workers)
  - OTel multi-agent semantic conventions:
      • ``execute_task``         — task decomposition & event propagation
      • ``invoke_agent``         — agent-to-agent interaction
      • ``agent_planning``       — internal routing / planning steps
      • ``agent.state.management`` — conversation context tracking
  - OTel tool conventions with standard attribute names:
      • ``tool.call.arguments``  — arguments passed to tool
      • ``tool.call.results``    — results returned by tool
  - Evaluation events attached to spans
  - User-feedback trace event

Trace hierarchy produced:
  orchestration_session
    └── provision_agents
          ├── provision:router_agent
          ├── provision:stats_agent
          ├── provision:matchup_agent
          └── provision:general_agent
    └── conversation_turn_{N}
          ├── task_decomposition                     ← task planning boundary
          │     └── intent_classification            ← router classifies intent
          │           └── route_intent:router_agent  ← router agent interaction
          │                 └── agent_execution      ← actual LLM run
          ├── delegate:stats_agent                   ← agent-to-agent delegation
          │     └── agent_execution                  ← actual LLM run
          │           └── tool_execution_batch       ← tool call processing
          │                 └── tool:get_team_stats  ← individual tool call
          │                 └── tool:get_player_info ← individual tool call
          ├── context_tracking                       ← conversation state
          └── evaluation                             ← quality assessment event

Agents:
  router_agent   — classifies user intent into: stats_lookup | matchup_analysis | general_question
  stats_agent    — has function tools (get_team_stats, get_player_info)
  matchup_agent  — has function tools (get_head_to_head)
  general_agent  — no tools, plain LLM response

Required environment variables:
  AZURE_AI_PROJECT_ENDPOINT  - Azure AI Foundry project endpoint
  APPLICATION_INSIGHTS_CONNECTION_STRING - (optional) enables cloud export
"""

import os
import json
import time
import uuid
from pathlib import Path
from datetime import datetime
from azure.ai.agents import AgentsClient
from azure.identity import DefaultAzureCredential
from azure.ai.agents.models import (
    FunctionTool,
    ListSortOrder,
    RequiredFunctionToolCall,
    ToolOutput,
)
from opentelemetry import trace
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from dotenv import load_dotenv

# ── Tracing imports ───────────────────────────────────────────────────────────
from tracing import ConsoleSpanExporter
from foundry_multi_agent_e2e_tracing import init_e2e_tracing, set_session_baggage


load_dotenv()

# Separate config file so agent IDs don't collide with other demos
AGENT_CONFIG_FILE = Path(__file__).parent / ".agent_config_multi.json"

# ── Bootstrap tracing — MUST happen before any Azure AI client is created ─────
connection_string = init_e2e_tracing(
    app_name="sports-multi-agent",
    namespace="foundry-analytics",
    version="1.0.0",
    agent_count=4,
    orchestration_pattern="router-worker",
)


# ──────────────────────────────────────────────────────────────────────────────
# Function-tool implementations (callable mock functions)
# ──────────────────────────────────────────────────────────────────────────────
# The SDK's FunctionTool introspects __name__, type-hints, and :param
# docstrings to auto-generate the JSON tool schemas for the LLM.
# ────────────────────────────────────────────────────────────────────────────── ─────────────────────────────────────────────────────

def get_team_stats(team_name: str, season: str = "2024") -> dict:
    """Get season statistics for a sports team.

    :param team_name: The team name, e.g. 'Dallas Cowboys'
    :param season: Season year, e.g. '2024'
    """
    return {
        "team": team_name,
        "season": season,
        "wins": 11,
        "losses": 6,
        "points_scored": 407,
        "points_allowed": 328,
        "division_rank": 2,
        "playoff_appearance": True,
    }


def get_player_info(player_name: str) -> dict:
    """Look up career information for a player by name.

    :param player_name: Full name of the player
    """
    return {
        "name": player_name,
        "position": "Quarterback",
        "team": "Mock City Legends",
        "career_years": 12,
        "career_touchdowns": 287,
        "career_yards": 42350,
        "pro_bowl_selections": 5,
        "championships": 2,
    }


def get_head_to_head(team_a: str, team_b: str) -> dict:
    """Get the head-to-head record between two teams.

    :param team_a: First team name
    :param team_b: Second team name
    """
    return {
        "team_a": team_a,
        "team_b": team_b,
        "all_time_record": {"team_a_wins": 34, "team_b_wins": 28, "ties": 2},
        "last_5_meetings": [
            f"{team_a} 27 - 24 {team_b}",
            f"{team_b} 31 - 17 {team_a}",
            f"{team_a} 20 - 20 {team_b} (OT)",
            f"{team_a} 35 - 28 {team_b}",
            f"{team_b} 14 - 10 {team_a}",
        ],
    }


# Map names → callable implementations (across all worker agents)
TOOL_IMPLEMENTATIONS = {
    "get_team_stats": get_team_stats,
    "get_player_info": get_player_info,
    "get_head_to_head": get_head_to_head,
}


def _build_function_tool(functions: set) -> FunctionTool:
    """Build a FunctionTool from a set of callable functions."""
    return FunctionTool(functions=functions)


# ──────────────────────────────────────────────────────────────────────────────
# Agent creation & management
# ──────────────────────────────────────────────────────────────────────────────

# Intent labels used by the router agent — keep in sync with AGENTS dict below
VALID_INTENTS = ["stats_lookup", "matchup_analysis", "general_question"]

# Agent definitions: (name, instructions, tool_functions_dict_or_None)
AGENT_SPECS = {
    "router": {
        "name": "router_agent",
        "model": "gpt-4.1",
        "instructions": (
            "You are a routing agent. Classify the user request into EXACTLY ONE of:\n"
            "  - stats_lookup       (team stats, player info, season data)\n"
            "  - matchup_analysis   (head-to-head, comparisons between teams)\n"
            "  - general_question   (opinions, history, rules, anything else)\n\n"
            "Reply with ONLY the intent label, nothing else."
        ),
        "tools": None,
    },
    "stats": {
        "name": "stats_agent",
        "model": "gpt-4.1",
        "instructions": (
            "You are a sports data specialist. ALWAYS use your tools to look up "
            "data before answering. Use get_team_stats and get_player_info to "
            "retrieve facts. Cite the data from tool results. Be concise."
        ),
        "tools": {get_team_stats, get_player_info},
    },
    "matchup": {
        "name": "matchup_agent",
        "model": "gpt-4.1",
        "instructions": (
            "You are a sports matchup analyst. ALWAYS use your get_head_to_head "
            "tool to retrieve matchup data before answering. Analyze trends from "
            "the data. Be concise."
        ),
        "tools": {get_head_to_head},
    },
    "general": {
        "name": "general_agent",
        "model": "gpt-4.1",
        "instructions": (
            "You are a knowledgeable sports commentator. Answer general sports "
            "questions with insight and enthusiasm. Be concise."
        ),
        "tools": None,
    },
}

# Map intent → worker agent key
INTENT_TO_WORKER = {
    "stats_lookup": "stats",
    "matchup_analysis": "matchup",
    "general_question": "general",
}


def _load_agent_config() -> dict:
    if AGENT_CONFIG_FILE.exists():
        with open(AGENT_CONFIG_FILE) as f:
            return json.load(f)
    return {}


def _save_agent_config(config: dict):
    with open(AGENT_CONFIG_FILE, "w") as f:
        json.dump(config, f, indent=2)


def create_all_agents(client: AgentsClient) -> dict[str, str]:
    """
    Create (or reuse) all agents defined in AGENT_SPECS.

    Returns a dict mapping agent key → agent ID.

    Uses ``invoke_agent`` semantic convention spans for each agent creation,
    with ``tool_definitions`` attribute per the OTel multi-agent spec.
    """
    tracer = trace.get_tracer(__name__)
    config = _load_agent_config()
    agent_ids: dict[str, str] = {}

    for key, spec in AGENT_SPECS.items():
        with tracer.start_as_current_span(f"provision:{spec['name']}") as span:
            span.set_attribute("gen_ai.agent.name", spec["name"])

            # Try to reuse existing agent
            cached_id = config.get(f"agent_id_{key}")
            if cached_id:
                span.set_attribute("agent.reused", True)
                try:
                    agent = client.get_agent(cached_id)
                    agent_ids[key] = cached_id
                    print(f"  ✓ Reusing {key} agent: {cached_id}")

                    # Record tool_definitions attribute (OTel semconv)
                    if spec["tools"]:
                        span.set_attribute(
                            "tool_definitions",
                            json.dumps([f.__name__ for f in spec["tools"]]),
                        )
                    continue
                except Exception:
                    print(f"  ⚠ Saved {key} agent not found — recreating...")

            span.set_attribute("agent.reused", False)

            # Build tools if specified
            tool_defs = []
            if spec["tools"]:
                ft = _build_function_tool(spec["tools"])
                tool_defs = ft.definitions
                span.set_attribute(
                    "tool_definitions",
                    json.dumps([f.__name__ for f in spec["tools"]]),
                )

            agent = client.create_agent(
                model=spec["model"],
                name=spec["name"],
                instructions=spec["instructions"],
                tools=tool_defs if tool_defs else None,
            )
            agent_ids[key] = agent.id
            config[f"agent_id_{key}"] = agent.id
            print(f"  ✓ Created {key} agent: {agent.id}")

    _save_agent_config(config)
    return agent_ids


# ──────────────────────────────────────────────────────────────────────────────
# Tool-call execution (with OTel semconv attributes)
# ──────────────────────────────────────────────────────────────────────────────

def execute_tool_call(function_name: str, arguments: dict) -> str:
    """
    Execute a single tool call locally and return the JSON result.

    Uses OTel semantic convention attribute names:
      - ``tool.call.arguments``  (matches the multi-agent observability spec)
      - ``tool.call.results``    (matches the multi-agent observability spec)
    """
    tracer = trace.get_tracer(__name__)

    with tracer.start_as_current_span(f"tool:{function_name}") as span:
        span.set_attribute("tool.function.name", function_name)
        span.set_attribute("tool.call.arguments", json.dumps(arguments))

        impl = TOOL_IMPLEMENTATIONS.get(function_name)
        if impl is None:
            error = f"Unknown function: {function_name}"
            span.set_status(trace.Status(trace.StatusCode.ERROR, error))
            return json.dumps({"error": error})

        try:
            result = impl(**arguments)
            result_json = json.dumps(result)
            span.set_attribute("tool.call.results", result_json)
            span.set_attribute("tool.status", "success")
            return result_json
        except Exception as exc:
            error = f"Tool execution failed: {exc}"
            span.set_status(trace.Status(trace.StatusCode.ERROR, error))
            return json.dumps({"error": error})


def handle_tool_calls(client: AgentsClient, thread_id: str, run) -> object:
    """
    Process all pending tool calls, submit outputs, and continue.

    Creates a parent ``handle_tool_calls`` span with child spans for each
    individual function execution.
    """
    tracer = trace.get_tracer(__name__)

    while run.status == "requires_action":
        tool_calls = run.required_action.submit_tool_outputs.tool_calls

        with tracer.start_as_current_span("tool_execution_batch") as span:
            span.set_attribute("tool_calls.count", len(tool_calls))
            span.set_attribute("thread.id", thread_id)
            span.set_attribute("run.id", run.id)

            tool_outputs = []
            for tc in tool_calls:
                if isinstance(tc, RequiredFunctionToolCall):
                    fn_name = tc.function.name
                    try:
                        args = json.loads(tc.function.arguments)
                    except json.JSONDecodeError:
                        args = {}

                    print(f"      🛠️  Executing tool: {fn_name}({json.dumps(args)})")
                    result = execute_tool_call(fn_name, args)
                    tool_outputs.append(ToolOutput(tool_call_id=tc.id, output=result))

            span.set_attribute(
                "tool_calls.functions",
                json.dumps([
                    tc.function.name
                    for tc in tool_calls
                    if isinstance(tc, RequiredFunctionToolCall)
                ]),
            )

        # Submit outputs and poll for completion
        client.runs.submit_tool_outputs(
            thread_id=thread_id,
            run_id=run.id,
            tool_outputs=tool_outputs,
        )

        # Poll until the run leaves requires_action
        while True:
            time.sleep(1)
            run = client.runs.get(thread_id=thread_id, run_id=run.id)
            if run.status not in ("queued", "in_progress"):
                break

    return run


# ──────────────────────────────────────────────────────────────────────────────
# Multi-agent orchestration spans (OTel GenAI semantic conventions)
# ──────────────────────────────────────────────────────────────────────────────

def _run_agent(
    client: AgentsClient,
    agent_id: str,
    thread_id: str,
    user_message: str,
    *,
    handle_tools: bool = True,
) -> str:
    """
    Send a message to an agent and return its response.

    Handles tool calls if ``handle_tools`` is True and the agent requests them.
    """
    tracer = trace.get_tracer(__name__)

    # Add user message
    client.messages.create(
        thread_id=thread_id,
        role="user",
        content=user_message,
    )

    # Run agent — manual polling to enable custom tool-call tracing
    with tracer.start_as_current_span("agent_execution") as run_span:
        run_span.set_attribute("thread.id", thread_id)
        run_span.set_attribute("agent.id", agent_id)

        run = client.runs.create(
            thread_id=thread_id,
            agent_id=agent_id,
        )
        run_span.set_attribute("run.id", run.id)

        # Poll until terminal state, handling tool calls manually
        while run.status in ("queued", "in_progress", "requires_action"):
            if run.status == "requires_action" and handle_tools:
                run = handle_tool_calls(client, thread_id, run)
                if run.status == "failed":
                    error_msg = f"Run failed after tool calls: {run.last_error}"
                    run_span.set_status(trace.Status(trace.StatusCode.ERROR, error_msg))
                    run_span.set_attribute("run.status", run.status)
                    return error_msg
            else:
                time.sleep(1)
                run = client.runs.get(thread_id=thread_id, run_id=run.id)

        run_span.set_attribute("run.status", run.status)

        if run.status == "failed":
            error_msg = f"Run failed: {run.last_error}"
            run_span.set_status(trace.Status(trace.StatusCode.ERROR, error_msg))
            return error_msg

    # Retrieve response
    messages = client.messages.list(
        thread_id=thread_id,
        order=ListSortOrder.DESCENDING,
        limit=1,
    )
    for msg in messages:
        if msg.role == "assistant" and msg.text_messages:
            return msg.text_messages[-1].text.value
    return "No response received"


def classify_intent(
    client: AgentsClient,
    router_agent_id: str,
    question: str,
) -> str:
    """
    Use the router agent to classify intent.

    Wrapped in ``agent_planning`` span (OTel multi-agent semantic convention)
    which captures the internal routing decision.  The nested
    ``invoke_agent.router`` span traces the actual agent-to-agent interaction.
    """
    tracer = trace.get_tracer(__name__)

    with tracer.start_as_current_span("intent_classification") as plan_span:
        plan_span.set_attribute("planning.type", "intent_classification")
        plan_span.set_attribute("planning.valid_intents", json.dumps(VALID_INTENTS))

        with tracer.start_as_current_span("route_intent:router_agent") as agent_span:
            agent_span.set_attribute("gen_ai.agent.name", "router_agent")
            agent_span.set_attribute("gen_ai.agent.id", router_agent_id)

            # Router gets its own thread so its classification doesn't
            # pollute the worker's conversation history.
            router_thread = client.threads.create()
            intent_raw = _run_agent(
                client,
                router_agent_id,
                router_thread.id,
                question,
                handle_tools=False,
            )
            intent = intent_raw.strip().lower()

            agent_span.set_attribute("router.raw_response", intent_raw)
            agent_span.set_attribute("router.intent", intent)

        # Validate intent
        if intent not in VALID_INTENTS:
            plan_span.set_attribute("planning.fallback", True)
            print(f"      ⚠ Unknown intent '{intent}', falling back to general_question")
            intent = "general_question"
        else:
            plan_span.set_attribute("planning.fallback", False)

        plan_span.set_attribute("planning.resolved_intent", intent)
        return intent


def delegate_to_worker(
    client: AgentsClient,
    agent_ids: dict[str, str],
    intent: str,
    question: str,
    thread_id: str,
) -> str:
    """
    Delegate a question to the appropriate worker agent.

    Wrapped in ``invoke_agent.<worker>`` span (OTel multi-agent semantic
    convention) to trace agent-to-agent interaction.
    """
    tracer = trace.get_tracer(__name__)
    worker_key = INTENT_TO_WORKER[intent]
    worker_agent_id = agent_ids[worker_key]

    worker_name = AGENT_SPECS[worker_key]["name"]

    with tracer.start_as_current_span(f"delegate:{worker_name}") as span:
        t0 = time.perf_counter()
        span.set_attribute("gen_ai.agent.name", worker_name)
        span.set_attribute("gen_ai.agent.id", worker_agent_id)
        span.set_attribute("worker.intent", intent)

        # Record tool definitions available to this worker
        worker_spec = AGENT_SPECS[worker_key]
        if worker_spec["tools"]:
            span.set_attribute(
                "tool_definitions",
                json.dumps([f.__name__ for f in worker_spec["tools"]]),
            )

        try:
            response = _run_agent(
                client,
                worker_agent_id,
                thread_id,
                question,
                handle_tools=True,
            )
            span.set_attribute("worker.status", "completed")
            span.set_attribute("worker.response.length", len(response))
            return response
        except Exception as exc:
            span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
            span.record_exception(exc)
            span.set_attribute("worker.status", "failed")
            raise
        finally:
            span.set_attribute(
                "worker.latency_ms",
                round((time.perf_counter() - t0) * 1000, 2),
            )


def track_context(
    tracer,
    turn_number: int,
    thread_id: str,
    intent: str,
    context_history: list[dict],
) -> None:
    """
    Record conversation context in an ``agent.state.management`` span.

    This satisfies the OTel multi-agent semantic convention for tracking
    effective context and short/long-term memory management.
    """
    with tracer.start_as_current_span("context_tracking") as span:
        span.set_attribute("state.turn_number", turn_number)
        span.set_attribute("state.thread_id", thread_id)
        span.set_attribute("state.current_intent", intent)
        span.set_attribute("state.history_length", len(context_history))
        span.set_attribute("state.context_summary", json.dumps(context_history[-3:]))


def record_evaluation(span, response: str, intent: str, turn_number: int) -> None:
    """
    Attach an evaluation event to the current span.

    This satisfies the OTel multi-agent semantic convention for structured
    evaluation of agent performance and decision-making.

    Event attributes follow the spec: name, error.type, label.
    """
    eval_attributes = {
        "evaluation.name": f"turn_{turn_number}_quality",
        "evaluation.label": intent,
        "evaluation.response_length": len(response),
        "evaluation.has_content": len(response.strip()) > 0,
    }

    # Simple quality heuristic for demo purposes
    if len(response.strip()) < 10:
        eval_attributes["error.type"] = "low_quality_response"
        eval_attributes["evaluation.score"] = 0.2
    elif "error" in response.lower() or "failed" in response.lower():
        eval_attributes["error.type"] = "error_in_response"
        eval_attributes["evaluation.score"] = 0.3
    else:
        eval_attributes["evaluation.score"] = 0.9

    span.add_event("evaluation", attributes=eval_attributes)


def record_user_feedback(span, turn_number: int, feedback: str, rating: str) -> None:
    """
    Attach a user-feedback event to the current span.

    Follows the OTel semantic convention for user feedback on GenAI responses:
      - gen_ai.user_feedback.rating  (thumbs_up | thumbs_down)
      - gen_ai.user_feedback.comment
    """
    span.add_event(
        "gen_ai.user.feedback",
        attributes={
            "gen_ai.user_feedback.rating": rating,
            "gen_ai.user_feedback.comment": feedback,
            "turn.number": turn_number,
        },
    )


# ──────────────────────────────────────────────────────────────────────────────
# Conversation turn (multi-agent)
# ──────────────────────────────────────────────────────────────────────────────

def run_turn(
    client: AgentsClient,
    agent_ids: dict[str, str],
    thread_id: str,
    user_message: str,
    turn_number: int,
    context_history: list[dict],
) -> tuple[str, str]:
    """
    Execute a single multi-agent conversation turn.

    1. ``execute_task`` — task decomposition boundary
    2. ``agent_planning`` → ``invoke_agent.router`` — classify intent
    3. ``invoke_agent.<worker>`` — delegate to specialist
    4. ``agent.state.management`` — track context
    5. ``evaluation`` event — assess response quality

    Returns:
        (response_text, resolved_intent)
    """
    tracer = trace.get_tracer(__name__)

    with tracer.start_as_current_span(f"conversation_turn_{turn_number}") as turn_span:
        t0 = time.perf_counter()
        turn_span.set_attribute("turn.number", turn_number)
        turn_span.set_attribute("user.query", user_message)
        turn_span.set_attribute("thread.id", thread_id)

        try:
            # ── execute_task: task planning & event propagation ───────────────
            with tracer.start_as_current_span("task_decomposition") as task_span:
                task_span.set_attribute("task.description", "classify_and_delegate")
                task_span.set_attribute(
                    "task.available_agents",
                    json.dumps(list(INTENT_TO_WORKER.values())),
                )

                # Classify intent via router agent
                intent = classify_intent(
                    client,
                    agent_ids["router"],
                    user_message,
                )
                task_span.set_attribute("task.resolved_intent", intent)
                task_span.set_attribute(
                    "task.delegated_to",
                    INTENT_TO_WORKER[intent],
                )

            turn_span.set_attribute("orchestrator.intent", intent)

            # ── invoke_agent.<worker>: agent-to-agent delegation ──────────────
            response = delegate_to_worker(
                client,
                agent_ids,
                intent,
                user_message,
                thread_id,
            )

            # ── agent.state.management: track context ─────────────────────────
            context_history.append({
                "turn": turn_number,
                "intent": intent,
                "query_preview": user_message[:80],
            })
            track_context(tracer, turn_number, thread_id, intent, context_history)

            # ── evaluation: assess response quality ───────────────────────────
            record_evaluation(turn_span, response, intent, turn_number)

            # ── simulate user feedback on odd turns (demo purposes) ───────────
            if turn_number % 2 == 1:
                record_user_feedback(
                    turn_span,
                    turn_number,
                    "Good answer, very helpful!",
                    "thumbs_up",
                )

            turn_span.set_attribute("turn.status", "completed")
            turn_span.set_attribute("response.length", len(response))
            return response, intent

        except Exception as exc:
            turn_span.set_status(
                trace.Status(trace.StatusCode.ERROR, str(exc)),
            )
            turn_span.record_exception(exc)
            turn_span.set_attribute("turn.status", "failed")
            raise

        finally:
            turn_span.set_attribute(
                "turn.latency_ms",
                round((time.perf_counter() - t0) * 1000, 2),
            )


# ──────────────────────────────────────────────────────────────────────────────
# Demo
# ──────────────────────────────────────────────────────────────────────────────

def demo():
    """
    Multi-agent orchestrated demo exercising all OTel multi-agent
    semantic conventions.

    Turn 1 — stats_lookup:      triggers Stats worker + function tools
    Turn 2 — matchup_analysis:  triggers Matchup worker + function tools
    Turn 3 — general_question:  triggers General worker (no tools)
    Turn 4 — stats_lookup:      triggers Stats worker again (context continuity)
    """
    project_endpoint = os.environ.get("AZURE_AI_PROJECT_ENDPOINT")
    if not project_endpoint:
        print("❌  AZURE_AI_PROJECT_ENDPOINT is not set.")
        print("    Export it and re-run:")
        print("      $env:AZURE_AI_PROJECT_ENDPOINT='https://<your-project>.api.azureml.ms'")
        return

    # ── output directory for local markdown traces ────────────────────────────
    output_dir = Path(__file__).parent / "output" / "traces"
    output_dir.mkdir(parents=True, exist_ok=True)
    session_ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    md_path = output_dir / f"trace_multi_agent_{session_ts}.md"
    md_file = open(md_path, "w", encoding="utf-8")
    md_file.write("# Foundry Multi-Agent Tracing Demo\n\n")
    md_file.write(
        f"**Session Start:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    )
    md_file.write("**Agents:** router_agent, stats_agent, matchup_agent, general_agent\n\n---\n\n")

    # ── attach console exporter ───────────────────────────────────────────────
    console_exporter = ConsoleSpanExporter(markdown_file=md_file)
    trace.get_tracer_provider().add_span_processor(
        SimpleSpanProcessor(console_exporter),
    )

    print(f"📝 Trace output → {md_path}\n")

    tracer = trace.get_tracer(__name__)

    # CREATE AND AUTHENTICATE
    client = AgentsClient(
        endpoint=project_endpoint,
        credential=DefaultAzureCredential(),
    )

    session_id = str(uuid.uuid4())
    set_session_baggage(session_id)  # propagates session.id to ALL spans via Baggage

    with tracer.start_as_current_span("orchestration_session") as session_span:
        session_span.set_attribute("session.id", session_id)
        session_span.set_attribute("demo.type", "multi_agent_tracing")
        session_span.set_attribute(
            "session.agents",
            json.dumps(list(AGENT_SPECS.keys())),
        )

        # ── Create all agents ────────────────────────────────────────────────
        print("Creating agents...")
        agent_ids = create_all_agents(client)
        session_span.set_attribute("agent.ids", json.dumps(agent_ids))

        # Create a shared thread for worker interactions
        thread = client.threads.create()
        thread_id = thread.id
        session_span.set_attribute("thread.id", thread_id)
        print(f"\n✓ Shared worker thread: {thread_id}\n")

        # ── Questions designed to exercise different workers ──────────────────
        turns = [
            # Turn 1 → stats_lookup (Stats worker + get_team_stats/get_player_info)
            "What are the Dallas Cowboys' season stats? Also look up Dak Prescott's career info.",

            # Turn 2 → matchup_analysis (Matchup worker + get_head_to_head)
            "Compare the head-to-head record between the Dallas Cowboys and Philadelphia Eagles.",

            # Turn 3 → general_question (General worker, no tools)
            "Why is American football called 'football' when players mostly use their hands?",

            # Turn 4 → stats_lookup (Stats worker again, tests context continuity)
            "Now look up Patrick Mahomes' player info and compare him to Prescott.",
        ]

        context_history: list[dict] = []

        print("=" * 70)
        print("FOUNDRY MULTI-AGENT TRACING PATTERNS — Demo")
        print("=" * 70)

        for i, question in enumerate(turns, start=1):
            print(f"\n{'─' * 70}")
            print(f"[TURN {i}] User: {question}")
            md_file.write(f"\n### 👤 User (Turn {i})\n\n{question}\n\n")

            response, intent = run_turn(
                client,
                agent_ids,
                thread_id,
                question,
                turn_number=i,
                context_history=context_history,
            )

            worker = INTENT_TO_WORKER[intent]
            print(f"[TURN {i}] Routed to: {worker} (intent: {intent})")
            print(
                f"[TURN {i}] Agent: "
                f"{response[:400]}{'...' if len(response) > 400 else ''}"
            )
            md_file.write(f"\n**Routed to:** `{worker}` (intent: `{intent}`)\n\n")
            md_file.write(f"\n### 🤖 Agent (Turn {i})\n\n{response}\n\n")

        session_span.set_attribute("total.turns", len(turns))
        session_span.set_attribute(
            "session.intents",
            json.dumps([h["intent"] for h in context_history]),
        )

        # ── Summary ──────────────────────────────────────────────────────────
        md_file.write("\n---\n\n## Summary\n\n")
        md_file.write(f"- **Session ID:** {session_id}\n")
        md_file.write(f"- **Turns:** {len(turns)}\n")
        md_file.write(f"- **Thread ID:** {thread_id}\n")
        md_file.write(f"- **Agents:** {json.dumps(agent_ids, indent=2)}\n")
        md_file.write(
            f"- **Intents:** {', '.join(h['intent'] for h in context_history)}\n"
        )
        md_file.write(
            f"- **Session End:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        )

    md_file.close()
    print(f"\n📄 Trace saved to: {md_path}")
    if connection_string:
        print(
            "✅  Application Insights export complete (allow 2–5 min for ingestion)"
        )


if __name__ == "__main__":
    demo()