orchestrator.py

import logging
import uuid
import os
from pprint import pprint
from graph.graph import configure_graph
from galileo import GalileoLogger

logging = logging.getLogger(__name__)

def main():
    # Initialize Galileo logger at the top level
    logger = GalileoLogger(
        project=os.getenv("GALILEO_PROJECT", "mason-trading-stocks"),
        log_stream=os.getenv("GALILEO_LOG_STREAM", "default")
    )
    
    # Start session once for the entire workflow
    session_id = f"trading-session-{uuid.uuid4()}"
    try:
        logger.start_session(name="stock-trading-session", external_id=session_id)
        print(f"Started Galileo session: {session_id}")
    except Exception as e:
        print(f"Galileo session error (continuing): {e}")
        session_id = None
    
    run_id = str(uuid.uuid4())
    
    try:
        # Configure and execute graph
        graph = configure_graph()
        
        g = graph.invoke({
            "id": run_id,
            "messages": [],
            "galileo_logger": logger  # Pass logger to nodes
        })
        
        print("\n\n Workflow completed: \n\n")
        pprint(g)
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        logging.error(f"Error: {str(e)}")
    
    # Clean up session
    if session_id:
        try:
            logger.flush()
        except Exception as e:
            print(f"Galileo session cleanup error: {e}")

if __name__ == "__main__":
    main()