import os
from mistralai.client import Mistral

client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])

# Start a workflow
execution = client.workflows.execute_workflow(
    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,
):
    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.broker_sequence + 1  # advance past this event
            process(event)
        break  # Completed normally
    except ConnectionError:
        await asyncio.sleep(1)
        # Loop will resume from last_seq

@workflows.workflow.define(name="parent-workflow")
class ParentWorkflow:
    @workflows.workflow.entrypoint
    async def run(self, params):
        # Child workflow events will have parent_workflow_exec_id set to this execution
        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 def resilient_consume(client, exec_id):
    max_retries = 10
    last_seq = 0

    last_seq = 0
    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.broker_sequence + 1  # advance past this event
                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(
    workflow_exec_id=exec_id, stream="token"
):
    process(event)

# Bad: Filter in code (wastes bandwidth)
async for event in client.workflows.events.get_stream_events(
    workflow_exec_id=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)