Spaces:
Sleeping
Sleeping
File size: 5,319 Bytes
285025c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
import os
from langgraph.graph import StateGraph, START, END
from IPython.display import Image, display
# Assuming 'agents' module and its contents (AgentState, nodes, routes) are available.
# For a runnable example, you'd need to define these or mock them.
# Example placeholders if 'agents.py' isn't provided:
class AgentState:
"""A placeholder for AgentState."""
pass
def analyze_intent_node(state):
"""Placeholder for analyze_intent_node."""
print("Analyzing intent...")
# In a real scenario, this would determine the next step
# For demonstration, let's simulate routing to visualization
return {"next_step": "visualization"}
def generate_visualization_node(state):
"""Placeholder for generate_visualization_node."""
print("Generating visualization code...")
# Simulate success
return {"code_generated": True}
def generate_transformation_node(state):
"""Placeholder for generate_transformation_node."""
print("Generating transformation code...")
return {"code_generated": True}
def generate_statistical_node(state):
"""Placeholder for generate_statistical_node."""
print("Generating statistical code...")
return {"code_generated": True}
def execute_code_node(state):
"""Placeholder for execute_code_node."""
print("Executing code...")
return {"execution_successful": True}
def error_handler_node(state):
"""Placeholder for error_handler_node."""
print("Handling error...")
return {}
def route_based_on_intent(state):
"""Placeholder for route_based_on_intent."""
# In a real app, this would use state to determine the route
if state.get("next_step") == "visualization":
return "visualization"
elif state.get("next_step") == "transformation":
return "transformation"
elif state.get("next_step") == "statistical":
return "statistical"
return "error"
def route_to_execution(state):
"""Placeholder for route_to_execution."""
# In a real app, this would check if code generation was successful
if state.get("code_generated"):
return "execute"
return "error"
def create_visualization():
"""Create and save a visualization of the agent workflow."""
# Create the state graph
workflow = StateGraph(AgentState)
# Add nodes
workflow.add_node("analyze_intent", analyze_intent_node)
workflow.add_node("visualization", generate_visualization_node)
workflow.add_node("transformation", generate_transformation_node)
workflow.add_node("statistical", generate_statistical_node)
workflow.add_node("execute", execute_code_node)
workflow.add_node("error_handler", error_handler_node)
# Add edges
workflow.add_edge(START, "analyze_intent")
# Conditional edges based on intent
workflow.add_conditional_edges(
"analyze_intent",
route_based_on_intent,
{
"visualization": "visualization",
"transformation": "transformation",
"statistical": "statistical",
"error": "error_handler"
}
)
# Route from generation nodes to execution
workflow.add_conditional_edges(
"visualization",
route_to_execution,
{
"execute": "execute",
"complete": END, # Added 'complete' to allow direct END from visualization if needed
"error": "error_handler"
}
)
workflow.add_conditional_edges(
"transformation",
route_to_execution,
{
"execute": "execute",
"complete": END,
"error": "error_handler"
}
)
workflow.add_conditional_edges(
"statistical",
route_to_execution,
{
"execute": "execute",
"complete": END,
"error": "error_handler"
}
)
# Final edges
workflow.add_edge("execute", END)
workflow.add_edge("error_handler", END)
# Create visualization directory if it doesn't exist
os.makedirs("visualizations", exist_ok=True)
# Generate and save the visualization
graph = workflow.compile()
try:
# Get the graph as a Mermaid diagram and draw it to PNG
# This requires 'mermaid-py' and potentially 'puppeteer' (for playwright backend)
png_data = graph.get_graph().draw_mermaid_png()
# Define the filename
filename = os.path.join("visualizations", "mermaid_graph.png")
# Save the PNG data to a file
with open(filename, "wb") as f:
f.write(png_data)
print(f"Image successfully saved as '{filename}'")
# Optionally, display the image after saving
display(Image(png_data))
except ImportError:
print("Please install 'mermaid-py' to generate PNG visualizations.")
print("You might also need to install a browser automation tool like 'playwright' for mermaid-py.")
except Exception as e:
print(f"An error occurred during visualization: {e}")
# This requires some extra dependencies and is optional
pass
if __name__ == "__main__":
create_visualization()
|