scripts/3.1_test_340b_traversal.py

import requests

def print_tree(node, prefix="", visited=None):
    """Recursive function to print a refined tree structure without duplicates."""
    if visited is None:
        visited = set()
    
    node_id = node.get("node_id")
    if not node_id:
        print(f"{prefix}(unknown node)")
        return

    # Avoid printing duplicate nodes by checking if node has been visited
    if node_id in visited:
        print(f"{prefix}(already listed) {node_id}")
        return
    visited.add(node_id)
    
    relationship_label = f"({node.get('relationship', '')})" if node.get("relationship") else ""
    print(f"{prefix}{node_id} {relationship_label}")
    
    children = node.get("descendants", [])
    for i, child in enumerate(children):
        new_prefix = f"{prefix}├── " if i < len(children) - 1 else f"{prefix}└── "
        print_tree(child, new_prefix, visited)

def inspect_and_print_relationships(node_id):
    """Inspect and display relationships for a specified node in both downward and upward directions."""
    # Downward traversal for child relationships
    response_down = requests.get(f"{base_url}/traverse_node?node_id={node_id}&direction=down")
    traversal_hierarchy_down = response_down.json().get("traversal_path", {})

    print(f"\nTraversal Response for {node_id} (Descendants):", traversal_hierarchy_down)  # Debugging line
    print(f"\nInspect Relationships for {node_id} (Descendants):")
    print_tree(traversal_hierarchy_down)

    # Upward traversal for parent relationships
    response_up = requests.get(f"{base_url}/traverse_node?node_id={node_id}&direction=up")
    traversal_hierarchy_up = response_up.json().get("traversal_path", {})

    print(f"\nTraversal Response for {node_id} (Ancestors):", traversal_hierarchy_up)  # Debugging line
    print(f"\nInspect Relationships for {node_id} (Ancestors):")
    print_tree(traversal_hierarchy_up)

# Base URL for API
base_url = "http://localhost:5000"

# Step 1: Load Graph (Specify the graph to load, e.g., PHSA/340B section)
print("\n--- Testing Graph Loading ---")
graph_data = {"graph_file": "graphs/PHSA/phsa_sec_340b.json"}
response = requests.post(f"{base_url}/load_graph", json=graph_data)
print("Load Graph Response:", response.json())

# Step 2: Test Upward and Downward Traversal for 340B Program
print("\n--- Testing Inspect Relationships for Node (340B Program) ---")
inspect_and_print_relationships("340B Program")