run_agent.py

from dotenv import load_dotenv
load_dotenv()

from agent.graph import build_graph

# Initial state — only set what we know at the start
initial_state = {
    'forecast_id':       '',
    'forecast_horizon':  48,
    'forecast_data':     [],
    'peak_forecast':     0.0,
    'classification':    '',
    'demand_response':   None,
    'procurement_flag':  None,
    'alert_payload':     None,
    'daily_briefing':    None,
    'decisions':         [],
    'notification_sent': False,
    'error':             None,
}

print("Building graph...")
graph = build_graph()

print("Running agent...")
result = graph.invoke(initial_state)

print(f"\nFinal classification: {result['classification']}")
print(f"Forecast ID: {result['forecast_id']}")
print(f"Peak forecast: {result['peak_forecast']:,.0f} MW")

if result['daily_briefing']:
    print(f"\nDaily Briefing:\n{result['daily_briefing']}")

if result['demand_response']:
    print(f"\nDemand Response:\n{result['demand_response']}")

if result['alert_payload']:
    print(f"\nAlert:\n{result['alert_payload']}")

if result['procurement_flag']:
    print(f"\nProcurement Flag:\n{result['procurement_flag']}")

if result['error']:
    print(f"\nError: {result['error']}")
    
# Test HIGH_DEMAND path
print("\n" + "="*60)
print("TESTING HIGH_DEMAND PATH")
print("="*60)

test_state = {
    **initial_state,
    'forecast_id':      'test-high-demand',
    'forecast_horizon': 48,
    'forecast_data':    [
        {
            'timestamp':    f'2018-08-03 {i:02d}:00:00',
            'predicted_mw': 160000.0,
            'lower_bound':  152000.0,
            'upper_bound':  168000.0,
        }
        for i in range(48)
    ],
    'peak_forecast': 160000.0,
}

# Run just classify + demand_response
from agent.nodes.classify import classify_forecast
from agent.nodes.demand_response import generate_demand_response

state = classify_forecast(test_state)
print(f"Classification: {state['classification']}")

state = generate_demand_response(state)
print(f"Demand Response: {state['demand_response']}")