CoolFace
Apppublic

syedkhizarrayaz/BM-AI-Analysis-And-Alert-Prioritization-Agent

sourceHugging Faceupdated 8mo agoView on Hugging Face
0likes
test_predict_api_to_excel.py261 linesDownload Raw Back to root
1"""2Script to test predicttransactionsjson API endpoint and save results to Excel3Reads alerts from JSON file, calls API, and saves to Excel with request and response columns4"""5 6import json7import requests8import pandas as pd9from typing import List, Dict, Any, Optional10import os11from datetime import datetime12from pathlib import Path13 14# Get script directory for relative paths15SCRIPT_DIR = Path(__file__).parent.absolute()16 17# Configuration - using relative paths18JSON_FILE_PATH = SCRIPT_DIR / "aml_alerts_5_test_predict.json"19OUTPUT_EXCEL_PATH = SCRIPT_DIR / "predict_api_test_results.xlsx"20 21# Common API ports to try22COMMON_PORTS = [8000, 8080, 5000, 3000, 8001]23 24 25def detect_api_base_url() -> str:26    """Automatically detect the API base URL by trying common ports"""27    print("Detecting API server...")28    29    # Try common localhost ports30    for port in COMMON_PORTS:31        base_url = f"http://localhost:{port}"32        test_endpoint = f"{base_url}/api/ai-service/predictalertpriority"33        34        try:35            # Try a quick HEAD request to check if server is up36            response = requests.head(f"{base_url}/", timeout=2)37            if response.status_code in [200, 404, 405]:  # 404/405 means server is up but endpoint might not exist38                print(f"✓ Found API server at: {base_url}")39                return base_url40        except requests.exceptions.ConnectionError:41            continue42        except requests.exceptions.Timeout:43            continue44        except Exception:45            continue46    47    # If nothing found, default to 800048    default_url = "http://localhost:8000"49    print(f"⚠ Could not auto-detect API server, using default: {default_url}")50    print("   Make sure your API server is running!")51    return default_url52 53 54# Auto-detect API base URL55API_BASE_URL = detect_api_base_url()56API_ENDPOINT = f"{API_BASE_URL}/api/ai-service/predictalertpriority"57 58 59def load_alerts_from_json(file_path: Path) -> List[Dict[str, Any]]:60    """Load alerts from JSON file"""61    try:62        # Convert to string if Path object63        file_path_str = str(file_path)64        if not os.path.exists(file_path_str):65            print(f"✗ JSON file not found: {file_path_str}")66            return []67        68        with open(file_path_str, 'r', encoding='utf-8') as f:69            data = json.load(f)70        if isinstance(data, list):71            return data72        else:73            return [data]74    except Exception as e:75        print(f"✗ Error loading JSON file: {e}")76        return []77 78 79def call_api(alert_data: Dict[str, Any], api_endpoint: str) -> Dict[str, Any]:80    """Call the API endpoint with alert data"""81    try:82        alert_id = alert_data.get('AlertID', 'Unknown')83        print(f"Calling API for AlertID: {alert_id}...")84        print(f"   Endpoint: {api_endpoint}")85        86        response = requests.post(87            api_endpoint,88            json=alert_data,89            headers={"Content-Type": "application/json"},90            timeout=300  # 5 minutes timeout91        )92        93        if response.status_code == 200:94            result = response.json()95            print(f"✓ Success for AlertID: {alert_id}")96            return result97        else:98            print(f"✗ API Error {response.status_code} for AlertID: {alert_id}")99            try:100                error_detail = response.json()101                error_msg = error_detail.get('detail', f"API returned status {response.status_code}")102            except:103                error_msg = f"API returned status {response.status_code}"104            105            return {106                "status": "error",107                "message": error_msg,108                "data": []109            }110    except requests.exceptions.ConnectionError:111        print(f"✗ Connection error for AlertID: {alert_data.get('AlertID', 'Unknown')}")112        print(f"   Make sure the API server is running at {API_BASE_URL}")113        return {114            "status": "error",115            "message": f"Could not connect to API server at {API_BASE_URL}",116            "data": []117        }118    except requests.exceptions.Timeout:119        print(f"✗ Timeout for AlertID: {alert_data.get('AlertID', 'Unknown')}")120        return {121            "status": "error",122            "message": "Request timeout (exceeded 5 minutes)",123            "data": []124        }125    except Exception as e:126        print(f"✗ Error calling API: {e}")127        return {128            "status": "error",129            "message": str(e),130            "data": []131        }132 133 134def process_alerts_and_save_to_excel():135    """Main function to process alerts and save to Excel"""136    print("=" * 60)137    print("API Test Script - Predict Transactions JSON")138    print("=" * 60)139    140    # Load alerts from JSON141    print(f"\n1. Loading alerts from: {JSON_FILE_PATH}")142    alerts = load_alerts_from_json(JSON_FILE_PATH)143    144    if not alerts:145        print("✗ No alerts found in JSON file!")146        return147    148    print(f"   ✓ Found {len(alerts)} alert(s) to process\n")149    150    # Process each alert151    all_results = []152    153    for idx, alert in enumerate(alerts, 1):154        print(f"\n{'='*60}")155        print(f"Processing Alert {idx}/{len(alerts)}")156        print(f"{'='*60}")157        158        # Call API159        api_response = call_api(alert, API_ENDPOINT)160        161        # Extract prediction data from response162        prediction_data = {}163        if api_response.get("status") == 200 and api_response.get("data"):164            # Get the first (and likely only) prediction result165            pred_data = api_response["data"][0] if api_response["data"] else {}166            prediction_data = {167                "Prediction": pred_data.get("Prediction", "N/A"),168                "STRScenario": pred_data.get("STRScenario", "N/A"),169                "FocusColumnValue_Response": pred_data.get("FocusColumnValue", "N/A")170            }171        elif api_response.get("status") == "error":172            prediction_data = {173                "Prediction": f"Error: {api_response.get('message', 'Unknown error')}",174                "STRScenario": "N/A",175                "FocusColumnValue_Response": "N/A"176            }177        else:178            prediction_data = {179                "Prediction": "No prediction returned",180                "STRScenario": "N/A",181                "FocusColumnValue_Response": "N/A"182            }183        184        # Combine request data with response185        result_row = alert.copy()186        result_row.update(prediction_data)187        188        # Add metadata from API response if available189        if api_response.get("status") == 200:190            result_row["API Status"] = api_response.get("status", "N/A")191            result_row["API Message"] = api_response.get("message", "N/A")192        193        all_results.append(result_row)194    195    # Create DataFrame196    print(f"\n{'='*60}")197    print("Creating Excel file...")198    print(f"{'='*60}")199    200    df = pd.DataFrame(all_results)201    202    # Reorder columns to put predictions near the end (but before metadata)203    columns = list(df.columns)204    prediction_cols = ["Prediction", "STRScenario", "FocusColumnValue_Response"]205    metadata_cols = ["API Status", "API Message"]206    207    # Remove prediction and metadata cols from main list208    for col in prediction_cols + metadata_cols:209        if col in columns:210            columns.remove(col)211    212    # Insert prediction cols before metadata213    insert_pos = len(columns)214    for col in metadata_cols:215        if col in columns:216            insert_pos = columns.index(col)217            break218    219    for col in reversed(prediction_cols):220        if col in df.columns:221            columns.insert(insert_pos, col)222    223    df = df[columns]224    225    # Save to Excel226    try:227        output_path_str = str(OUTPUT_EXCEL_PATH)228        df.to_excel(output_path_str, index=False, engine='openpyxl')229        print(f"\n✓ Successfully saved results to: {output_path_str}")230        print(f"  Total rows: {len(df)}")231        print(f"  Total columns: {len(df.columns)}")232        print(f"\nColumns in Excel:")233        for i, col in enumerate(df.columns, 1):234            print(f"  {i}. {col}")235    except ImportError:236        print(f"\n⚠ openpyxl not installed, saving as CSV instead...")237        print("   Install with: pip install openpyxl")238        csv_path = str(OUTPUT_EXCEL_PATH).replace('.xlsx', '.csv')239        df.to_csv(csv_path, index=False)240        print(f"✓ Saved as CSV: {csv_path}")241    except Exception as e:242        print(f"\n✗ Error saving to Excel: {e}")243        print("Trying to save as CSV instead...")244        csv_path = str(OUTPUT_EXCEL_PATH).replace('.xlsx', '.csv')245        df.to_csv(csv_path, index=False)246        print(f"✓ Saved as CSV: {csv_path}")247 248 249if __name__ == "__main__":250    try:251        process_alerts_and_save_to_excel()252        print("\n" + "=" * 60)253        print("Script completed successfully!")254        print("=" * 60)255    except KeyboardInterrupt:256        print("\n\nScript interrupted by user")257    except Exception as e:258        print(f"\n✗ Fatal error: {e}")259        import traceback260        traceback.print_exc()261