Optimize my travel booking agent's prompts to improve routing accuracy

# Create the skill directory
mkdir -p .claude/skills/training-apo-agents

# Copy the skill file (or create from template)
# The skill lives at: .claude/skills/training-apo-agents/SKILL.md

# Install Agent-Lightning with APO support
uv add "agentlightning[apo]"

# You'll also need these for LangChain agents
uv add langchain langchain-openai langgraph langgraph-swarm

# Required environment variables
export OPENROUTER_API_KEY=your_key_here    # For APO algorithm
export LANGSMITH_API_KEY=your_key_here     # Optional: for dataset loading
export LANGSMITH_TRACING=true              # Enable tracing

I have a multi-agent travel system with triage, routing, and specialist agents.
The triage agent often routes to the wrong specialist. Optimize the prompts
to improve routing accuracy. My evaluation dataset is in LangSmith as 
"multi_agent_evals_dataset".

I'll optimize your travel system prompts using APO. Here's the plan:

Configuration:
- Algorithm: APO with beam search
- Hardware: Local CPU with 4 parallel runners
- Beam width: 4 prompts, 3 optimization rounds
- Estimated time: ~20-30 minutes

Components to create:
1. Prompt templates for all 6 agents
2. Factory function to rebuild swarm with new prompts
3. Reward function with routing accuracy as primary metric
4. LitAgent wrapper for APO training

Should I proceed with the smoke test first?

# Explicit reward weights - adjust based on your priorities
REWARD_WEIGHTS = {
    "term_matching": 0.30,      # Output contains expected terms
    "tool_calls": 0.25,         # Correct tools were called
    "agent_routing": 0.30,      # Correct agent sequence
    "efficiency": 0.15,         # Step count within bounds
}

uv run train_multi_agent.py --config small

uv run train_multi_agent.py --config full

Training complete!
Best optimized prompt template:
----------------------------------------
You are a triage agent for a travel company. Your ONLY job is to
analyze user intent and IMMEDIATELY route to specialists using
transfer tools...
----------------------------------------
Optimized prompts saved to: optimized_prompts_20260116_223615.py

                    ┌─────────────┐
                    │ Seed Prompt │
                    └──────┬──────┘
                           │
              ┌────────────┼────────────┐
              ▼            ▼            ▼
         ┌────────┐  ┌────────┐  ┌────────┐
         │Variant1│  │Variant2│  │Variant3│  ← Branch Factor
         └───┬────┘  └───┬────┘  └───┬────┘
             │           │           │
         [Evaluate]  [Evaluate]  [Evaluate]
             │           │           │
             └─────┬─────┴─────┬─────┘
                   ▼           ▼
              ┌────────┐  ┌────────┐
              │ Top 2  │  │ Top 2  │          ← Beam Width
              └───┬────┘  └───┬────┘
                  │           │
                  └─────┬─────┘
                        ▼
                 [Next Round]                  ← Beam Rounds

# training/prompts.py
import agentlightning as agl

TRIAGE_PROMPT = """You are a triage agent for a travel company. Your job is to 
route users to the correct specialist.
Based on user intent, IMMEDIATELY call the appropriate transfer tool:
- For flight or hotel BOOKINGS → call transfer_to_booking_router_agent
- For flight delay/cancellation CLAIMS → call transfer_to_claims_agent
- For general INFO/FAQ questions → call transfer_to_faq_agent
IMPORTANT: You MUST use the transfer tools. Do NOT respond with text."""

def get_initial_resources() -> dict[str, agl.PromptTemplate]:
    return {
        "triage_prompt": agl.PromptTemplate(
            template=TRIAGE_PROMPT,
            engine="f-string"
        ),
        # Add more prompts for multi-agent systems...
    }

# training/swarm_factory.py
from langchain.agents import create_agent
from langgraph_swarm import create_handoff_tool, create_swarm
from langgraph.checkpoint.memory import MemorySaver
import agentlightning as agl

def create_swarm_with_prompts(resources: agl.NamedResources):
    """Rebuild swarm with optimized prompts."""
    llm = get_llm()
    
    # Extract prompts from resources
    triage_prompt = resources["triage_prompt"].format()
    router_prompt = resources["router_prompt"].format()
    
    # Create agents with current prompts
    triage_agent = create_agent(
        name="triage_agent",
        model=llm,
        system_prompt=triage_prompt,
        tools=[
            create_handoff_tool("booking_router_agent", "For bookings"),
            create_handoff_tool("claims_agent", "For claims"),
            create_handoff_tool("faq_agent", "For questions"),
        ],
    )
    
    # ... create other agents ...
    
    workflow = create_swarm(
        [triage_agent, ...],
        default_active_agent="triage_agent"
    )
    return workflow.compile(checkpointer=MemorySaver())

# training/reward.py
from langchain_core.messages import AIMessage

REWARD_WEIGHTS = {
    "term_matching": 0.30,      # Output contains expected terms
    "tool_calls": 0.25,         # Correct tools were called
    "agent_routing": 0.30,      # Correct agent sequence
    "efficiency": 0.15,         # Step count within bounds
}

def extract_trajectory_info(result: dict) -> dict:
    """Extract what happened during agent execution."""
    messages = result.get("messages", [])
    
    # Get final output
    final_output = ""
    for msg in reversed(messages):
        if isinstance(msg, AIMessage) and msg.content:
            final_output = msg.content
            break
    
    # Get tool calls
    tool_calls = []
    for msg in messages:
        if isinstance(msg, AIMessage) and msg.tool_calls:
            tool_calls.extend([tc["name"] for tc in msg.tool_calls])
    
    # Get agent routing sequence
    agent_sequence = ["Triage Agent"]
    agent_mapping = {
        "transfer_to_flight_booking_agent": "Flight Agent",
        "transfer_to_claims_agent": "Claims Agent",
        # ... etc
    }
    for msg in messages:
        if isinstance(msg, AIMessage) and msg.tool_calls:
            for tc in msg.tool_calls:
                if tc["name"] in agent_mapping:
                    agent_sequence.append(agent_mapping[tc["name"]])
    
    return {
        "final_output": final_output,
        "tool_calls": tool_calls,
        "agent_sequence": agent_sequence,
        "step_count": sum(1 for m in messages if isinstance(m, AIMessage)),
    }

def calculate_reward(outputs: dict, reference: dict) -> float:
    """Calculate weighted composite reward."""
    rewards = {}
    
    # 1. Term matching
    expected_terms = reference.get("expected_terms", [])
    if expected_terms:
        final_output = outputs.get("final_output", "").lower()
        matched = sum(1 for t in expected_terms if t.lower() in final_output)
        rewards["term_matching"] = matched / len(expected_terms)
    
    # 2. Tool calls
    expected_tools = reference.get("expected_tool_calls", [])
    if expected_tools:
        actual_tools = outputs.get("tool_calls", [])
        matched = sum(1 for t in expected_tools if t in actual_tools)
        rewards["tool_calls"] = matched / len(expected_tools)
    
    # 3. Agent routing (for multi-agent)
    expected_agents = reference.get("expected_agent_sequence", [])
    if expected_agents:
        agent_sequence = outputs.get("agent_sequence", [])
        # Subsequence matching
        exp_idx = 0
        for agent in agent_sequence:
            if exp_idx < len(expected_agents) and agent == expected_agents[exp_idx]:
                exp_idx += 1
        rewards["agent_routing"] = exp_idx / len(expected_agents)
    
    # 4. Efficiency
    step_count = outputs.get("step_count", 0)
    min_steps, max_steps = reference.get("min_steps", 1), reference.get("max_steps", 10)
    if min_steps <= step_count <= max_steps:
        rewards["efficiency"] = 1.0
    else:
        rewards["efficiency"] = max(0.0, 1.0 - abs(step_count - max_steps) / max_steps)
    
    # Weighted sum
    total = sum(rewards.get(k, 0) * v for k, v in REWARD_WEIGHTS.items())
    return total / sum(REWARD_WEIGHTS.values())

# training/lit_agent.py
import uuid
import agentlightning as agl
from typing import TypedDict

class TravelTask(TypedDict):
    prompt: str
    expected_terms: list[str]
    expected_tool_calls: list[str]
    expected_agent_sequence: list[str]
    min_steps: int
    max_steps: int

class TravelSwarmLitAgent(agl.LitAgent[TravelTask]):
    def __init__(self):
        super().__init__()
        self._initial_resources = get_initial_resources()
    
    def rollout(
        self, 
        task: TravelTask, 
        resources: agl.NamedResources, 
        rollout: agl.Rollout
    ) -> float:
        # Merge optimized prompts with fixed prompts
        full_resources = dict(self._initial_resources)
        full_resources.update(resources)
        
        # Build and run agent
        app = create_swarm_with_prompts(full_resources)
        thread_id = f"training-{uuid.uuid4()}"
        
        try:
            result = app.invoke(
                {"messages": [{"role": "user", "content": task["prompt"]}]},
                config={"configurable": {"thread_id": thread_id}}
            )
        except Exception as e:
            print(f"Agent error: {e}")
            return 0.0
        
        outputs = extract_trajectory_info(result)
        return calculate_reward(outputs, task)

# train_multi_agent.py

# ===== MUST BE FIRST - before any other imports =====
import multiprocessing
try:
    multiprocessing.set_start_method("fork", force=True)
except RuntimeError:
    pass
# ====================================================

import os
import argparse
from dotenv import load_dotenv
from openai import AsyncOpenAI
import agentlightning as agl

from training import get_initial_resources, TravelSwarmLitAgent
from training.config import get_config

load_dotenv()
os.environ["LANGSMITH_TRACING"] = "true"

def main(config_name: str = "full"):
    config = get_config(config_name)
    
    # Configure OpenRouter client
    client = AsyncOpenAI(
        api_key=os.getenv("OPENROUTER_API_KEY"),
        base_url="https://openrouter.ai/api/v1"
    )
    
    # Configure APO
    algo = agl.APO(
        client,
        beam_width=config.beam_width,
        branch_factor=config.branch_factor,
        beam_rounds=config.beam_rounds,
    )
    
    # Load dataset
    train_data, val_data = load_dataset(config.dataset_name)
    
    # Configure Trainer
    trainer = agl.Trainer(
        algorithm=algo,
        n_runners=config.n_runners,
        initial_resources=get_initial_resources(),
        adapter=agl.TraceToMessages(),
    )
    
    # Run training
    trainer.fit(
        agent=TravelSwarmLitAgent(),
        train_dataset=train_data,
        val_dataset=val_data,
    )
    
    # Get optimized prompt
    best_prompt = algo.get_best_prompt()
    print(f"Optimized: {best_prompt.template}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", choices=["small", "full"], default="full")
    args = parser.parse_args()
    main(args.config)

SMALL_CONFIG = TrainingConfig(
    beam_width=1,
    branch_factor=1,
    beam_rounds=1,
    n_runners=1,
    train_split=0.5,
)

FULL_CONFIG = TrainingConfig(
    beam_width=4,
    branch_factor=4,
    beam_rounds=3,
    n_runners=4,
    train_split=0.7,
)

def calculate_reward(outputs: dict, reference: dict) -> float:
    rewards = {}
    
    # ... calculate individual rewards ...
    
    # Debug logging
    print(f"  term_matching: {rewards.get('term_matching', 0):.2f}")
    print(f"  tool_calls: {rewards.get('tool_calls', 0):.2f}")
    print(f"  agent_routing: {rewards.get('agent_routing', 0):.2f}")
    print(f"  efficiency: {rewards.get('efficiency', 0):.2f}")
    
    total = ...
    print(f"  TOTAL: {total:.2f}")
    return total

                          User Query
                              │
                              ▼
                    ┌─────────────────┐
                    │  Triage Agent   │ ← Prompt optimized
                    └────────┬────────┘
                             │
           ┌─────────────────┼─────────────────┐
           ▼                 ▼                 ▼
    ┌──────────────┐  ┌──────────────┐  ┌──────────────┐
    │Booking Router│  │ Claims Agent │  │  FAQ Agent   │
    └──────┬───────┘  └──────────────┘  └──────────────┘
           │                    All prompts optimized
     ┌─────┴─────┐
     ▼           ▼
┌──────────┐ ┌──────────┐
│  Flight  │ │  Hotel   │
│  Agent   │ │  Agent   │
└──────────┘ └──────────┘

# In reward.py
expected_agents = reference.get("expected_agent_sequence", [])
# e.g., ["Triage Agent", "Booking Router Agent", "Flight Agent"]

actual_sequence = outputs.get("agent_sequence", [])
# e.g., ["Triage Agent", "Booking Router Agent", "Flight Agent"]

# Check subsequence matching (order matters)
exp_idx = 0
for agent in actual_sequence:
    if exp_idx < len(expected_agents) and agent == expected_agents[exp_idx]:
        exp_idx += 1
routing_score = exp_idx / len(expected_agents)

task = {
    "prompt": "I want to book a flight to Paris",
    "expected_terms": ["flight", "paris"],
    "expected_tool_calls": ["search_flights"],
    "expected_agent_sequence": [
        "Triage Agent", 
        "Booking Router Agent", 
        "Flight Booking Agent"
    ],
    "min_steps": 3,
    "max_steps": 8,
}

TRIAGE_PROMPT = """You are a triage agent for a travel company. 
Your job is to route users to the correct specialist.
Based on user intent, IMMEDIATELY call the appropriate transfer tool:
- For flight or hotel BOOKINGS → call transfer_to_booking_router_agent
- For claims → call transfer_to_claims_agent
- For general questions → call transfer_to_faq_agent"""

OPTIMIZED_TRIAGE_PROMPT = """You are a triage agent for TravelCo. 
Your ONLY job is to analyze user intent and IMMEDIATELY route.

ROUTING RULES (apply strictly in order):
1. BOOKING keywords (book, reserve, find, search + flight/hotel) 
   → transfer_to_booking_router_agent
2. CLAIM keywords (delay, cancel, refund, compensation, claim)
   → transfer_to_claims_agent  
3. ALL OTHER queries (info, question, policy, how, what, when)
   → transfer_to_faq_agent

CRITICAL: Call transfer tool immediately. Never respond with text.
If uncertain between booking and claim, prefer claims_agent."""