from mistralai.workflows.client import get_mistral_client

async with get_mistral_client(
    server_url="https://api.mistral.ai",
    api_key=os.environ["MISTRAL_API_KEY"],
) as client:
    # Start a workflow
    execution = await client.workflows.execute_workflow_async(
        workflow_identifier="my-workflow",
        input=params,
    )

    # Stream all events for this execution
    async for event in client.workflows.events.get_stream_events(
        workflow_exec_id=execution.execution_id,
    ):
        print(f"[{event.stream}] {event.data}")

async for event in client.workflows.events.get_stream_events(
    workflow_exec_id=execution_id,
    stream="token",  # Only token events
):
    print(event.data)

async for event in client.workflows.events.get_stream_events(
    workflow_exec_id=execution_id,
):
    if event.data is None:
        continue

    event_type = event.data.event_type

    if event_type == "WORKFLOW_EXECUTION_COMPLETED":
        print("Done!")
        break
    elif event_type in ("WORKFLOW_EXECUTION_FAILED", "WORKFLOW_EXECUTION_CANCELED"):
        print(f"Ended: {event_type}")
        break
    else:
        record_event(event)

last_seq = 0
while True:
    try:
        async for event in client.workflows.events.get_stream_events(
            workflow_exec_id=execution_id,
            start_seq=last_seq,
        ):
            last_seq = event.sequence
            process(event)
        break  # Completed normally
    except ConnectionError:
        await asyncio.sleep(1)
        # Loop will resume from last_seq

# Parent workflow
@workflows.workflow.define(name="parent-workflow")
class ParentWorkflow:
    @workflows.workflow.entrypoint
    async def run(self, params):
        # Sub-workflow events will have parent_workflow_exec_id set
        result = await workflows.workflow.execute_workflow(
            ChildWorkflow,
            params
        )
        return result

# Get events from parent AND all children
async for event in client.workflows.events.get_stream_events(
    root_workflow_exec_id=parent_execution_id,
):
    if event.workflow_context.parent_workflow_exec_id:
        print(f"[child] {event.data}")
    else:
        print(f"[parent] {event.data}")

async for event in client.workflows.events.get_stream_events(
    root_workflow_exec_id="root123",
):
    process(event)

async def resilient_consume(client, exec_id):
    max_retries = 10

    for attempt in range(max_retries):
        try:
            async for event in client.workflows.events.get_stream_events(
                workflow_exec_id=exec_id,
                start_seq=last_seq,
            ):
                last_seq = event.sequence
                yield event

                if is_terminal_status(event):
                    return
        except ConnectionError:
            await asyncio.sleep(min(2 ** attempt, 30))

# Good: Filter at source
async for event in client.workflows.events.get_stream_events(exec_id, stream="token"):
    process(event)

# Bad: Filter in code (wastes bandwidth)
async for event in client.workflows.events.get_stream_events(exec_id):
    if event.stream == "token":
        process(event)

async for event in client.workflows.events.get_stream_events(
    workflow_exec_id=exec_id, scope="activity"
):
    if event.data is not None:
        # Only activity events, no workflow-level events
        process(event.data)