Strategies

from ziet import Agent, ManualApprovalStrategy, PickActionsStrategy, HandOffStrategy, input

# Agent with manual approval
@ManualApprovalStrategy(signee=input.user_email)
@Agent(
    id="booking_agent",
    name="BookingAgent",
    instructions="Book flight after user approval",
    actions=["book_flight", "charge_payment"]
)
class BookingAgent:
    pass

# Agent that picks best actions dynamically
@PickActionsStrategy(picks=2, allow_repeats=False)
@Agent(
    id="research_agent",
    name="ResearchAgent",
    instructions="Research topic using best available sources",
    actions=["search_google", "search_news", "scrape_web"]
)
class ResearchAgent:
    pass

# Agent that hands off to another agent
@HandOffStrategy(agents=["review_agent"])
@Agent(
    id="draft_agent",
    name="DraftAgent",
    instructions="Create draft and hand off for review",
    actions=["generate_draft"]
)
class DraftAgent:
    pass

@StrategyDecorator(params)
@Agent(
    id="my_agent",
    name="MyAgent",
    instructions="...",
    actions=["action1", "action2"]
)
class MyAgent:
    pass

from ziet import Agent, ManualApprovalStrategy, input

@ManualApprovalStrategy(
    signee=input.user_email,  # Who to ask for approval
    method="email"             # How to request: "email" | "slack" | "sms"
)
@Agent(
    id="deletion_agent",
    name="DeletionAgent",
    instructions="""
    Delete user account and all associated data.
    This is a sensitive operation requiring approval.
    """,
    actions=["delete_account", "remove_data", "send_notification"]
)
class DeletionAgent:
    pass

from ziet import Agent, WaitStrategy

@WaitStrategy(
    memory=["user_approval", "payment_confirmed"],
    sleep=10,       # Check every 10 seconds
    timeout=3600    # Wait max 1 hour
)
@Agent(
    id="fulfillment_agent",
    name="FulfillmentAgent",
    instructions="""
    Process and fulfill order after approval and payment are confirmed.
    Retrieve 'user_approval' and 'payment_confirmed' from memory.
    """,
    actions=["process_order", "ship_order", "send_tracking"]
)
class FulfillmentAgent:
    pass

from ziet import Agent, CritiqueStrategy

@CritiqueStrategy(
    loops=2,           # Number of critique iterations
    model="gpt-4o"     # Model for critique
)
@Agent(
    id="content_agent",
    name="ContentAgent",
    instructions="""
    Generate high-quality marketing email for the product and audience.
    The CritiqueStrategy will refine your output through multiple iterations.
    Store final result in memory with key 'final_email'.
    """,
    actions=["generate_draft", "analyze_tone"],
    model="gpt-4o-mini"
)
class ContentAgent:
    pass

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

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


@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

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

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

@ManualApprovalStrategy(signee=input.user_email)
@WaitStrategy(memory=["payment_confirmed"])
@CritiqueStrategy(loops=1)
@Agent(
    id="invoicing_agent",
    name="InvoicingAgent",
    instructions="""
    Generate and send invoice with multiple safeguards:
    1. Wait for 'payment_confirmed' in memory (WaitStrategy)
    2. Generate invoice with AI refinement (CritiqueStrategy)
    3. Require user approval before sending (ManualApprovalStrategy)
    4. Send invoice to customer
    """,
    actions=["generate_invoice", "send_invoice_email"],
    model="gpt-4o-mini"
)
class InvoicingAgent:
    pass

from ziet import Agent, ManualApprovalStrategy, input

@ManualApprovalStrategy(signee=input.user_email)
@Agent(
    id="approval_agent",
    name="ApprovalAgent",
    instructions="""
    Execute sensitive workflow with user approval.
    The signee for approval is dynamically set from the request.
    """,
    actions=["sensitive_action", "log_activity"]
)
class ApprovalAgent:
    pass

# When invoked:
# {
#   "user_email": "manager@company.com",  # Used as signee
#   "data": {...}
# }

from ziet import Agent, ManualApprovalStrategy, input

@ManualApprovalStrategy(signee=input.approver_email)
@Agent(
    id="deployment_agent",
    name="DeploymentAgent",
    instructions="Deploy application to production environment",
    actions=["run_tests", "deploy_code", "verify_deployment"]
)
class DeploymentAgent:
    pass

from ziet import Agent, HandOffStrategy

# Agent 1: Processes payment
@HandOffStrategy(agents=["shipping_agent"])
@Agent(
    id="payment_agent",
    name="PaymentAgent",
    instructions="""
    Process payment and store result in memory with key 'payment_confirmed'.
    Then hand off to shipping_agent.
    """,
    actions=["charge_payment", "verify_transaction"]
)
class PaymentAgent:
    pass

# Agent 2: Waits for payment, then ships
@WaitStrategy(memory=["payment_confirmed"])
@Agent(
    id="shipping_agent",
    name="ShippingAgent",
    instructions="""
    Wait for 'payment_confirmed' in memory, then ship the order.
    """,
    actions=["prepare_shipment", "create_label", "dispatch"]
)
class ShippingAgent:
    pass

from ziet import Agent, CritiqueStrategy

@CritiqueStrategy(loops=3, model="gpt-4o")
@Agent(
    id="writer_agent",
    name="WriterAgent",
    instructions="""
    Write a high-quality blog post on the given topic for the target audience.
    The CritiqueStrategy will refine your output through 3 iterations.
    Store final post in memory with key 'blog_post'.
    """,
    actions=["research_topic", "generate_outline", "write_content"],
    model="gpt-4o-mini"
)
class WriterAgent:
    pass

from ziet import Agent, ManualApprovalStrategy, CritiqueStrategy, input

@ManualApprovalStrategy(signee=input.manager_email)
@CritiqueStrategy(loops=2)
@Agent(
    id="communication_agent",
    name="CommunicationAgent",
    instructions="""
    Generate customer communication email.
    1. CritiqueStrategy refines the content (2 loops)
    2. ManualApprovalStrategy requires manager approval
    3. After approval, send the email
    """,
    actions=["generate_email", "send_email"],
    model="gpt-4o"
)
class CommunicationAgent:
    pass

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) -> None:
    query = f"flights from {origin} to {dest} on {date}"
    results = google.search(query, num_results=10)
    memory.add(key="google_flight_results", value=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) -> None:
    # Search Kayak, Expedia, etc.
    results = google.search(f"{origin} to {dest} kayak expedia")
    memory.add(key="comparison_site_results", value=results)

@Action(
    id="scrape_airline",
    name="Scrape Airline",
    description="Scrape airline website",
    timeout=60
)
def scrape_airline(url: str) -> None:
    data = apify.scrape(url)
    scraped_data = {
        "url": url,
        "content": data.get("text", "")[:1000],
        "prices": data.get("extracted", {})
    }
    memory.add(key=f"scraped_{url}", value=scraped_data)

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

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

@Action(
    id="send_confirmation",
    name="Send Confirmation",
    description="Send confirmation email",
    timeout=20
)
def send_confirmation(email: str) -> None:
    # 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
    )
    memory.add(key="confirmation_sent", value={"email": email, "sent": 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. Combine all results and store 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: stores results in 'google_flight_results'
    - search_flight_sites: stores results in 'comparison_site_results'
    - scrape_airline: stores results in 'scraped_{url}'
    
    After actions complete, retrieve from memory and combine into 'all_flight_options'.
    """,
    actions=["search_google_flights", "search_flight_sites", "scrape_airline"],
    model="gpt-4o-mini"
)
class FlightSearcherAgent:
    pass


# ============================================
# 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 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 (retrieves 'recommended_flight' from memory)
    8. Send confirmation email using send_confirmation action (retrieves 'final_booking' from memory)
    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.
    All actions store their results in memory for traceability.
    """,
    actions=["book_flight", "charge_payment", "send_confirmation"],
    model="gpt-4o-mini"
)
class FlightBookerAgent:
    pass


# ============================================
# 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