test1.py

# test1.py

import logging
import traceback
from datetime import datetime

logging.basicConfig(level=logging.INFO)

from src.state import VideoGenerationState
from src.nodes.consultation.consultation_pipeline import ConsultationPipelineNode


MOCK_USERS = [
    "mock_user_priya_001",
    "mock_user_rahul_002",
    "mock_user_anjali_003",
    "mock_user_kiran_004",
    "mock_user_deepak_005",
    "mock_user_arjun_007",
    "mock_user_vikram_009",
    "mock_user_neha_010",
    "mock_user_farhan_011",
    "mock_user_kavya_012",
    "mock_user_rohit_013",
    "mock_user_aisha_014",
    "mock_user_nitin_015",
    "mock_user_pooja_016",
    "mock_user_aditya_017",
    "mock_user_sana_018",
]


def generate_for_user(user_id: str, node: ConsultationPipelineNode):
    print("\n" + "=" * 80)
    print(f"Generating consultation video for: {user_id}")
    print("=" * 80)

    state = VideoGenerationState(
        session_id=user_id,
        topic="consultation",
        target_language="english",
        optional_params={
            "script_type": "consultation"
        }
    )

    result = node.generate_consultation_video(state)

    print(f"\n✓ SUCCESS: {user_id}")
    print(f"Status : {result.status}")
    print(f"Video  : {result.slide_video_path}")

    return result


def main():

    start_time = datetime.now()

    print("\n")
    print("=" * 100)
    print("CONSULTATION VIDEO BATCH GENERATION")
    print(f"Users to process: {len(MOCK_USERS)}")
    print("=" * 100)

    node = ConsultationPipelineNode()

    success = []
    failed = []

    for idx, user_id in enumerate(MOCK_USERS, start=1):

        print("\n")
        print(f"[{idx}/{len(MOCK_USERS)}] Processing {user_id}")

        try:
            result = generate_for_user(user_id, node)

            success.append({
                "user": user_id,
                "status": result.status,
                "video": result.slide_video_path
            })

        except Exception as e:

            print(f"\n✗ FAILED: {user_id}")
            print(str(e))

            traceback.print_exc()

            failed.append({
                "user": user_id,
                "error": str(e)
            })

    end_time = datetime.now()

    print("\n")
    print("=" * 100)
    print("BATCH GENERATION SUMMARY")
    print("=" * 100)

    print(f"\nTotal Users     : {len(MOCK_USERS)}")
    print(f"Successful      : {len(success)}")
    print(f"Failed          : {len(failed)}")
    print(f"Started         : {start_time}")
    print(f"Finished        : {end_time}")
    print(f"Duration        : {end_time - start_time}")

    print("\nSUCCESSFUL USERS")
    print("-" * 100)

    for item in success:
        print(f"{item['user']} -> {item['video']}")

    if failed:
        print("\nFAILED USERS")
        print("-" * 100)

        for item in failed:
            print(f"{item['user']} -> {item['error']}")

    print("\nDone.")


if __name__ == "__main__":
    main()