Policy

import asyncio
import json
from dedalus_labs import AsyncDedalus, DedalusRunner
from dedalus_labs.utils.streaming import stream_async
from dotenv import load_dotenv

load_dotenv()

def on_tool(evt: dict) -> None:
    print("[policy tool evt]", json.dumps(evt))

def add(a: int, b: int) -> int:
    print(f"[tool:add] a={a} b={b}")
    return a + b

def mul(x: int, y: int) -> int:
    print(f"[tool:mul] x={x} y={y}")
    return x * y

def policy(ctx: dict) -> dict:
    step = ctx.get("step", 1)
    print(f"[policy] step={step}")
    pol: dict = {}

    if step == 3:
        pol.update({"message_prepend": [{"role": "system", "content": "You must speak like a pirate."}]})
    
    # Cap total steps for safety
    pol.setdefault("max_steps", 4)
    return pol

async def main() -> None:
    client = AsyncDedalus()
    runner = DedalusRunner(client)

    prompt = (
        "Step 1) Add 7 and 8. "
        "Step 2) Multiply the previous result by 3. "
        "Step 3) Search the web for the final number and summarize the first result title. "
    )

    result = runner.run(
        input=prompt,
        model="gpt-4.1",
        tools=[add, mul],
        mcp_servers=["dedalus-labs/brave-search-mcp"],
        stream=True,
        on_tool_event=on_tool,
        policy=policy
    )

    await stream_async(result)    

if __name__ == "__main__":
    asyncio.run(main())