main.py

"""
DS-STAR: Multi-Agent System for Data Science Tasks

This is the main entry point for the refactored DS-STAR system.
All agents are modularized and can be tested independently.

Usage:
    python main_new.py

Or customize:
    from src.graph import run_ds_star
    from src.config import get_llm

    llm = get_llm(provider="google", model="gemini-1.5-flash")
    result = run_ds_star("Your question here", llm, max_iterations=20)
"""

import sys

from src.config import DEFAULT_CONFIG, get_llm
from src.graph import run_ds_star


def main():
    """
    Main execution function for DS-STAR.
    """
    # Configuration
    query = "What percentage of transactions use credit cards?"
    max_iterations = DEFAULT_CONFIG["max_iterations"]
    provider = DEFAULT_CONFIG["provider"]
    model = DEFAULT_CONFIG["model"]

    print("Initializing DS-STAR Multi-Agent System...")
    print(f"Provider: {provider}")
    print(f"Model: {model}")
    print()

    try:
        # Initialize LLM
        llm = get_llm(
            provider=provider, model=model, temperature=DEFAULT_CONFIG["temperature"]
        )

        # Run DS-STAR workflow
        final_state = run_ds_star(
            query=query,
            llm=llm,
            max_iterations=max_iterations,
            thread_id="ds-star-main-session",
        )

        if final_state:
            print("\n✅ Workflow completed successfully!")
            return 0
        else:
            print("\n❌ Workflow failed!")
            return 1

    except Exception as e:
        print(f"\n❌ Fatal error: {str(e)}")
        import traceback

        traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())