Graphify / timeline_generator.py
ZahirJS's picture
Update timeline_generator.py
18eb0e2 verified
raw
history blame
7.14 kB
import graphviz
import json
from tempfile import NamedTemporaryFile
import os
def generate_timeline_diagram(json_input: str, output_format: str) -> str:
"""
Generates a serpentine timeline diagram from JSON input.
Args:
json_input (str): A JSON string describing the timeline structure.
It must follow the Expected JSON Format Example below.
Expected JSON Format Example:
{
"title": "AI Development Timeline",
"events_per_row": 4,
"events": [
{
"id": "event_1",
"label": "Machine Learning Foundations",
"date": "1950-1960",
"description": "Early neural networks and perceptrons"
},
{
"id": "event_2",
"label": "Expert Systems Era",
"date": "1970-1980",
"description": "Rule-based AI systems"
},
{
"id": "event_3",
"label": "Neural Network Revival",
"date": "1980-1990",
"description": "Backpropagation algorithm"
}
]
}
Returns:
str: The filepath to the generated PNG image file.
"""
try:
if not json_input.strip():
return "Error: Empty input"
data = json.loads(json_input)
if 'events' not in data:
raise ValueError("Missing required field: events")
dot = graphviz.Digraph(
name='Timeline',
format='png',
graph_attr={
'rankdir': 'TB', # Top-to-Bottom for better control
'splines': 'ortho', # Straight lines with 90-degree bends
'bgcolor': 'white', # White background
'pad': '0.8', # Padding around the graph
'nodesep': '3.0', # Increased spacing between nodes horizontally
'ranksep': '2.5' # Increased spacing between ranks vertically
}
)
base_color = '#19191a' # Hardcoded base color
title = data.get('title', '')
events = data.get('events', [])
events_per_row = data.get('events_per_row', 4) # Default to 4 events per row
if not events:
raise ValueError("Timeline must contain at least one event")
# Add title node if provided
if title:
dot.node(
'title',
title,
shape='plaintext',
fontsize='18',
fontweight='bold',
fontcolor=base_color,
pos="6,2!" # Centered at top (adjust x-coordinate based on timeline width)
)
total_events = len(events)
# Create all event nodes first
for i, event in enumerate(events):
event_id = event.get('id', f'event_{i}')
event_label = event.get('label', f'Event {i+1}')
event_date = event.get('date', '')
event_description = event.get('description', '')
# Create full label with date and description
if event_date and event_description:
full_label = f"{event_date}\\n{event_label}\\n{event_description}"
elif event_date:
full_label = f"{event_date}\\n{event_label}"
elif event_description:
full_label = f"{event_label}\\n{event_description}"
else:
full_label = event_label
# Calculate color for current depth using same logic as graph_generator_utils
lightening_factor = 0.12
current_depth = i # Use event position as depth for color progression
# Convert base_color hex to RGB for interpolation
if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
base_color_safe = '#19191a' # Fallback to default dark if invalid
else:
base_color_safe = base_color
base_r = int(base_color_safe[1:3], 16)
base_g = int(base_color_safe[3:5], 16)
base_b = int(base_color_safe[5:7], 16)
# Calculate current node color by blending towards white
current_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
current_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
current_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
# Clamp values to 255 to stay within valid RGB range
current_r = min(255, current_r)
current_g = min(255, current_g)
current_b = min(255, current_b)
node_color = f'#{current_r:02x}{current_g:02x}{current_b:02x}'
# Font color: white for dark nodes, black for very light nodes for readability
font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
# Font size adjusts based on position, ensuring a minimum size
font_size = max(9, 14 - (current_depth * 1))
# Calculate position for serpentine layout
row = i // events_per_row
col = i % events_per_row
# For odd rows, reverse the column position to create serpentine effect
if row % 2 == 1:
visual_col = events_per_row - 1 - col
else:
visual_col = col
# Add the event node with position attributes for layout
dot.node(
event_id,
full_label,
shape='box',
style='filled,rounded',
fillcolor=node_color,
fontcolor=font_color,
fontsize=str(font_size),
width='2.5',
height='1.2',
pos=f"{visual_col * 4.5},{-row * 3}!" # Increased spacing for serpentine layout
)
# Connect events in chronological order (1→2→3→4...)
for i in range(len(events) - 1):
current_event_id = events[i].get('id', f'event_{i}')
next_event_id = events[i + 1].get('id', f'event_{i + 1}')
dot.edge(
current_event_id,
next_event_id,
color='#4a4a4a', # Dark gray for lines (consistent with other diagrams)
arrowsize='0.8',
penwidth='2'
)
# No need to connect title to events - it stays at the top independently
# Set the layout engine to handle fixed positions
dot.engine = 'neato'
with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
dot.render(tmp.name, format=output_format, cleanup=True)
return f"{tmp.name}.{output_format}"
except json.JSONDecodeError:
return "Error: Invalid JSON format"
except Exception as e:
return f"Error: {str(e)}"