from alayaflow.api import Flow


def main():
    # Initialize flow api singleton

    flow = Flow()
    flow.init({
        "alayahub_url": "http://xxxx",
    })

    # Load workflows

    workflow_id = 'test_chat'
    workflow_version = '1.0.0'
    init_args = {
        "alayamem_url": "http://10.16.70.46:5555",
    }
    flow.load_workflow(workflow_id, workflow_version, init_args)

    # Execute workflow

    input_data = {
        "user_id": "test_user",
        "session_id": "test_session",
        "messages": [{"role": "user", "content": "你是谁"}],
    }
    user_config = {
        # You can set workflow runtime config here
        # Such as API key, temperature, etc., as long as they are defined in the workflow
    }
    print("Start executing workflow...")
    for event in flow.exec_workflow(workflow_id, workflow_version, input_data, user_config):
        if "error" in event:
            print(f"\n[Error] {event['error']}")
            break
        kind = event["event"]
        if kind == "on_chat_model_stream":
            content = event["data"]["chunk"]['content']
            if content:
                # 打印每一个 Token，实现打字机效果
                print(content, end="", flush=True)
    print()


if __name__ == "__main__":
    main()