Strategies

from ziet import Action, WaitStrategy, ManualApprovalStrategy, CritiqueStrategy, input

@ManualApprovalStrategy(signee=input.user_email)
@Action(description="Sensitive operation")
def delete_data(user_id: str) -> dict:
    # Requires approval before executing
    return {"deleted": user_id}

@WaitStrategy(memory=["approval_key"])
@Action(description="Wait for condition")
def process_after_approval() -> dict:
    # Waits until memory contains "approval_key"
    return {"status": "processed"}

@CritiqueStrategy(loops=2)
@Action(description="Generate email")
def generate_email(topic: str) -> dict:
    # LLM critiques and refines output
    return {"body": f"Email about {topic}"}

@StrategyDecorator(params)
@Action(description="...")
def my_action() -> dict:
    return result

from ziet import Action, ManualApprovalStrategy, input

@ManualApprovalStrategy(
    signee=input.user_email,  # Who to ask for approval
    method="email"             # How to request: "email" | "slack" | "sms"
)
@Action(description="Delete user account")
def delete_account(user_id: str) -> dict:
    # This won't execute until approved
    return database.delete_user(user_id)

from ziet import Action, WaitStrategy, memory

@WaitStrategy(
    memory=["user_approval", "payment_confirmed"],
    actions=["process_payment", "verify_address"],
    sleep=10,       # Check every 10 seconds
    timeout=3600    # Wait max 1 hour
)
@Action(description="Process order")
def process_order(order_id: str) -> dict:
    # Waits until:
    # 1. Both memory keys exist
    # 2. Both actions have completed
    
    approval = memory.get("user_approval")
    payment = memory.get("payment_confirmed")
    
    return {"order_id": order_id, "status": "processed"}

from ziet import Action, CritiqueStrategy

@CritiqueStrategy(
    loops=2,           # Number of critique iterations
    model="gpt-4o"     # Model for critique
)
@Action(description="Generate marketing email")
def generate_email(product: str, audience: str) -> dict:
    # Initial generation
    body = f"Check out our new {product} for {audience}!"
    
    return {"subject": "New Product", "body": body}

from ziet import Agent, PickActionsStrategy

@PickActionsStrategy(
    picks=2,              # Execute 2 actions
    allow_repeats=False   # Don't pick same action twice
)
@Agent(
    id="search_agent",
    name="SearchAgent",
    instructions="""
    Research the given topic by searching multiple sources.
    The PickActionsStrategy will intelligently select the best 2 actions
    from available options based on context.
    Available actions: search_google, search_news, scrape_page, etc.
    Store findings in memory for later use.
    """,
    actions=["search_google", "search_news", "scrape_page"]
)
class SearchAgent:
    pass  # Agent behavior defined by instructions and strategies

from ziet import Agent, HandOffStrategy

@HandOffStrategy(agents=["booking_agent", "refund_agent"])
@Agent(
    id="search_agent",
    name="SearchAgent",
    instructions="""
    Search for flights between origin and destination.
    Store results in durable memory with key 'flight_options'.
    Hand off to booking_agent when search is complete.
    Memory persists across the handoff.
    """,
    actions=["search_flights"]
)
class SearchAgent:
    pass  # Agent behavior defined by instructions and strategies


@HandOffStrategy(agents=["search_agent"])
@Agent(
    id="booking_agent",
    name="BookingAgent",
    instructions="""
    Retrieve flight options from memory (key: 'flight_options').
    Book the best flight option.
    Memory from previous agent (SearchAgent) is available.
    Store booking confirmation in memory.
    """,
    actions=["book_flight"]
)
class BookingAgent:
    pass  # Agent behavior defined by instructions and strategies

from ziet import Agent, ConditionalStrategy

@ConditionalStrategy(condition="input.user_type == 'premium'")
@Agent(
    id="premium_agent",
    name="PremiumAgent",
    instructions="Handle premium user requests with priority features",
    actions=["premium_action_1", "premium_action_2"]
)
class PremiumAgent:
    pass


@ConditionalStrategy(condition="memory.get('score') > 80")
@Agent(
    id="highscore_agent",
    name="HighScoreAgent",
    instructions="Process high-scoring results",
    actions=["advanced_analysis"]
)
class HighScoreAgent:
    pass


# String matching conditions
@ConditionalStrategy(condition="input.answer == 'blue'")
@Agent(
    id="blue_agent",
    name="BlueAgent",
    instructions="Handle blue answers",
    actions=["process_blue"]
)
class BlueAgent:
    pass


# Complex conditions
@ConditionalStrategy(
    condition="memory.get('payment_status') == 'confirmed' and input.amount < 1000"
)
@Agent(
    id="autoapprove_agent",
    name="AutoApproveAgent",
    instructions="Auto-approve small confirmed payments",
    actions=["approve_payment", "send_receipt"]
)
class AutoApproveAgent:
    pass

# Example 1: User type routing
@ConditionalStrategy(condition="input.user_type == 'enterprise'")
@Agent(
    id="enterprise_agent",
    name="EnterpriseAgent",
    instructions="Handle enterprise customers with dedicated support",
    actions=["priority_support", "custom_solutions"]
)
class EnterpriseAgent:
    pass


# Example 2: Score-based processing
@ConditionalStrategy(condition="memory.get('confidence_score') > 0.9")
@Agent(
    id="highconfidence_agent",
    name="HighConfidenceAgent",
    instructions="Auto-process high confidence results without review",
    actions=["auto_approve"]
)
class HighConfidenceAgent:
    pass


# Example 3: Multi-condition logic
@ConditionalStrategy(
    condition="input.region == 'US' and memory.get('inventory') > 0"
)
@Agent(
    id="fulfillment_agent",
    name="USFulfillmentAgent",
    instructions="Fulfill orders for US region with available inventory",
    actions=["create_shipment", "notify_customer"]
)
class USFulfillmentAgent:
    pass


# Example 4: String matching
@ConditionalStrategy(condition="input.language in ['en', 'es', 'fr']")
@Agent(
    id="multilingual_agent",
    name="MultilingualAgent",
    instructions="Process supported languages",
    actions=["translate", "respond"]
)
class MultilingualAgent:
    pass


# Example 5: Negation
@ConditionalStrategy(condition="not memory.get('already_processed')")
@Agent(
    id="firsttime_agent",
    name="FirstTimeProcessor",
    instructions="Only process if not already done",
    actions=["process_data", "mark_complete"]
)
class FirstTimeProcessor:
    pass

@ConditionalStrategy(condition="input.amount > 10000")
@ManualApprovalStrategy(signee=input.approver_email)
@Agent(
    id="largetransaction_agent",
    name="LargeTransactionAgent",
    instructions="Handle large transactions with approval",
    actions=["process_payment", "send_notification"]
)
class LargeTransactionAgent:
    pass
# Only executes if amount > 10000, then requires manual approval

@ManualApprovalStrategy(signee=input.user_email)
@WaitStrategy(memory=["payment_confirmed"])
@CritiqueStrategy(loops=1)
@Action(description="Generate and send invoice")
def send_invoice(customer_id: str, amount: int) -> dict:
    # 1. Wait for payment confirmation
    # 2. Generate invoice (with critique)
    # 3. Request manual approval
    # 4. Send invoice
    
    invoice = generate_invoice_content(customer_id, amount)
    send_email(customer_id, invoice)
    
    return {"status": "sent", "amount": amount}

from ziet import Action, ManualApprovalStrategy, input

@ManualApprovalStrategy(signee=input.user_email)
@Action(
    id="sensitive_action",
    name="Sensitive Action",
    description="Sensitive action requiring approval"
)
def sensitive_action(user_email: str, data: dict) -> dict:
    # signee is resolved from request["user_email"] at runtime
    return {"status": "done"}

# Agent using this action:
@Agent(
    id="approval_agent",
    name="ApprovalAgent",
    instructions="Execute sensitive action with user approval",
    actions=["sensitive_action"]
)
class ApprovalAgent:
    pass  # The signee parameter uses input.user_email from the request

@ManualApprovalStrategy(signee=input.approver_email)
@Action(description="Deploy to production")
def deploy_to_production(version: str) -> dict:
    return deploy(version)

# Action 1: Process payment
@Action(description="Process payment")
def process_payment(amount: int) -> dict:
    result = stripe.charge(amount)
    memory.add(result, key="payment_confirmed")
    return result

# Action 2: Wait for payment, then ship
@WaitStrategy(memory=["payment_confirmed"])
@Action(description="Ship order")
def ship_order(order_id: str) -> dict:
    return {"status": "shipped"}

@CritiqueStrategy(loops=3, model="gpt-4o")
@Action(description="Write blog post")
def write_blog_post(topic: str, audience: str) -> dict:
    # Initial draft
    content = openai.chat([{
        "role": "user",
        "content": f"Write a blog post about {topic} for {audience}"
    }])
    
    return {"content": content}

@ManualApprovalStrategy(signee=input.manager_email)
@CritiqueStrategy(loops=2)
@Action(description="Generate customer communication")
def send_customer_email(customer_id: str, message: str) -> dict:
    # 1. Generate email (with critique)
    # 2. Request manager approval
    # 3. Send email
    
    email = generate_email(message)
    send_email(customer_id, email)
    
    return {"status": "sent"}

from ziet import Action, Agent, memory, input
from ziet import PickActionsStrategy, CritiqueStrategy, HandOffStrategy, ManualApprovalStrategy
from ziet.integrations import google, apify, openai, stripe, sendgrid

# ============================================
# ACTIONS
# ============================================

@Action(
    id="search_google_flights",
    name="Search Google for Flights",
    description="Search Google for flights",
    timeout=30
)
def search_google_flights(origin: str, dest: str, date: str) -> list:
    query = f"flights from {origin} to {dest} on {date}"
    results = google.search(query, num_results=10)
    return results

@Action(
    id="search_flight_sites",
    name="Search Flight Sites",
    description="Search flight comparison sites",
    timeout=30
)
def search_flight_sites(origin: str, dest: str) -> list:
    # Search Kayak, Expedia, etc.
    results = google.search(f"{origin} to {dest} kayak expedia")
    return results

@Action(
    id="scrape_airline",
    name="Scrape Airline",
    description="Scrape airline website",
    timeout=60
)
def scrape_airline(url: str) -> dict:
    data = apify.scrape(url)
    return {
        "url": url,
        "content": data.get("text", "")[:1000],
        "prices": data.get("extracted", {})
    }

@Action(
    id="book_flight",
    name="Book Flight",
    description="Book flight",
    timeout=60
)
def book_flight(flight: dict, payment: dict) -> dict:
    # Book the flight
    booking = {
        "booking_id": f"BK{flight['id']}",
        "flight": flight,
        "status": "confirmed"
    }
    memory.add(key="final_booking", value=booking)
    return booking

@Action(
    id="charge_payment",
    name="Charge Payment",
    description="Charge payment",
    timeout=30
)
def charge_payment(amount: int, email: str) -> dict:
    payment = stripe.create_payment_intent(
        amount=amount * 100,
        currency="usd",
        metadata={"customer": email}
    )
    return payment

@Action(
    id="send_confirmation",
    name="Send Confirmation",
    description="Send confirmation email",
    timeout=20
)
def send_confirmation(email: str) -> bool:
    # Retrieve booking from memory
    booking = memory.get("final_booking")
    
    sendgrid.send(
        to=email,
        subject=f"Flight Booked - {booking['booking_id']}",
        body=f"<h1>Booking Confirmed</h1><p>Booking ID: {booking['booking_id']}</p>",
        html=True
    )
    return True


# ============================================
# AGENT 1: FLIGHT SEARCHER
# ============================================

@PickActionsStrategy(picks=2, allow_repeats=False)
@HandOffStrategy(agents=["reviewer_agent"])
@Agent(
    id="searcher_agent",
    name="FlightSearcher",
    description="Search for flight options using multiple strategies",
    instructions="""
    You are a flight search agent. Your goal is to:
    1. Search multiple sources for flights (Google, comparison sites)
    2. Use the PickActionsStrategy to select the 2 best search methods
    3. Gather at least 10-20 flight options
    4. Store initial request in memory with key 'initial_request'
    5. Store search results in memory with key 'all_flight_options'
    6. Hand off to reviewer_agent when search is complete
    
    Prioritize variety in search sources for comprehensive results.
    The PickActionsStrategy will intelligently select from available actions:
    - search_google_flights
    - search_flight_sites
    - scrape_airline
    """,
    actions=["search_google_flights", "search_flight_sites", "scrape_airline"],
    model="gpt-4o-mini"
)
class FlightSearcherAgent:
    pass  # Agent behavior is defined by instructions and strategies


# ============================================
# AGENT 2: FLIGHT REVIEWER
# ============================================

@CritiqueStrategy(loops=2, model="gpt-4o")
@HandOffStrategy(agents=["booker_agent"])
@Agent(
    id="reviewer_agent",
    name="FlightReviewer",
    description="Review and analyze flight options",
    instructions="""
    You are a flight analysis expert. Your task:
    1. Retrieve all flight options from memory using key 'all_flight_options'
    2. Retrieve initial request from memory using key 'initial_request'
    3. Analyze each option considering price, duration, stops, airline quality
    4. Recommend the top 3 options with detailed reasoning
    5. The CritiqueStrategy will refine your analysis through 2 loops
    6. Store the final analysis in memory with key 'flight_analysis'
    7. Store the top recommendation in memory with key 'recommended_flight'
    8. Hand off to booker_agent when analysis is complete
    
    Consider overall value, not just price. Factor in convenience, reliability, 
    and travel experience. The CritiqueStrategy will help you refine your analysis.
    """,
    actions=[],
    model="gpt-4o"
)
class FlightReviewerAgent:
    pass  # Agent behavior is defined by instructions and strategies


# ============================================
# AGENT 3: FLIGHT BOOKER
# ============================================

@ManualApprovalStrategy(signee=input.user_email, method="email")
@Agent(
    id="booker_agent",
    name="FlightBooker",
    description="Book flight after manual approval",
    instructions="""
    You are responsible for the final booking. Your workflow:
    1. Retrieve recommended flight from memory using key 'recommended_flight'
    2. Retrieve flight analysis from memory using key 'flight_analysis'
    3. Retrieve initial request from memory using key 'initial_request'
    4. The ManualApprovalStrategy will send approval request to user via email
    5. WAIT for user approval (execution pauses automatically)
    6. After approval: charge payment using charge_payment action
    7. Book the flight using book_flight action
    8. Send confirmation email using send_confirmation action
    9. Store final result in memory with key 'final_result'
    
    Only proceed with payment after explicit approval. The ManualApprovalStrategy
    handles the approval workflow automatically. If the user rejects, execution stops.
    """,
    actions=["book_flight", "charge_payment", "send_confirmation"],
    model="gpt-4o-mini"
)
class FlightBookerAgent:
    pass  # Agent behavior is defined by instructions and strategies


# ============================================
# EXECUTION
# ============================================

# Start the workflow by calling searcher_agent
# Memory persists across all agent handoffs
# Each agent is stateless but shares durable memory

"""
Execution Flow:

1. searcher_agent (PickActionsStrategy + HandOffStrategy)
   - Intelligently picks 2 search actions from available pool
   - Searches multiple sources (Google, flight sites, airline scraping)
   - Stores results in memory: 'all_flight_options', 'initial_request'
   - Hands off to reviewer_agent
   - Memory persists across handoff

2. reviewer_agent (CritiqueStrategy + HandOffStrategy)
   - Reads flight options from memory
   - Analyzes with GPT-4
   - CritiqueStrategy refines analysis (2 loops)
   - Stores refined analysis in memory: 'flight_analysis', 'recommended_flight'
   - Hands off to booker_agent
   - Memory persists across handoff

3. booker_agent (ManualApprovalStrategy)
   - Reads recommended flight from memory
   - ManualApprovalStrategy sends approval request via email
   - Execution pauses waiting for approval
   - After approval: charges payment, books flight, sends confirmation
   - Stores final result in memory: 'final_result'

Memory is durable throughout:
- If any step fails, memory persists
- Agents can retry from last successful step
- No data loss across handoffs
- Each agent is stateless but accesses shared memory
"""

# Deploy all agents
ziet deploy

# Start the workflow (calls searcher_agent first)
ziet run searcher_agent \
  --origin "SFO" \
  --dest "NYC" \
  --date "2024-12-25" \
  --user_email "user@example.com" \
  --max_price 600