File size: 4,080 Bytes
5a2d62e | 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 | import os
from datetime import datetime
from typing import Dict
import subprocess
class FileManager:
def __init__(self):
self.output_base = "output"
def create_output_dir(self, project_name: str, timestamp: str) -> str:
"""Create output directory structure for this run."""
dir_name = f"{project_name}_{timestamp}"
output_dir = os.path.join(self.output_base, dir_name)
# Create main directories
os.makedirs(output_dir, exist_ok=True)
os.makedirs(os.path.join(output_dir, "code"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "analysis"), exist_ok=True)
os.makedirs(os.path.join(output_dir, "images"), exist_ok=True)
return output_dir
def save_step_result(self, output_dir: str, step_name: str, content: str):
"""Save intermediate step result with clean PlantUML code."""
ext = ".pu" if "diagram" in step_name else ".md"
# Determine target directory based on file type
if ext == ".pu":
target_dir = os.path.join(output_dir, "code")
# Render PlantUML diagram
self._render_plantuml(content, output_dir, step_name)
else:
target_dir = os.path.join(output_dir, "analysis")
filename = f"{step_name}{ext}"
filepath = os.path.join(target_dir, filename)
print(f"Saving {filename} to {target_dir}") # Add this line for debugging
# Clean PlantUML code by removing markdown code block markers
if ext == ".pu":
content = self._clean_plantuml_content(content)
with open(filepath, "w") as f:
f.write(content)
def _render_plantuml(self, content: str, output_dir: str, step_name: str):
"""Render PlantUML diagram to PNG."""
# Clean the content first
content = self._clean_plantuml_content(content)
# Create temporary .pu file
temp_pu = os.path.join(output_dir, "code", f"{step_name}.pu")
with open(temp_pu, "w") as f:
f.write(content)
# Render to PNG
output_png = os.path.join(output_dir, "images", f"{step_name}.png")
try:
subprocess.run([
"plantuml",
"-tpng",
temp_pu,
"-o", "images"
], check=True)
except subprocess.CalledProcessError as e:
print(f"Error rendering PlantUML diagram: {e}")
except FileNotFoundError:
print("PlantUML not found. Please install PlantUML to generate images.")
def _clean_plantuml_content(self, content: str) -> str:
"""Remove triple quotes and plantuml markers from the content."""
# Remove '''plantuml and ''' markers
content = content.replace("'''plantuml", "").replace("'''", "")
content = content.replace("```plantuml", "").replace("```", "")
# Remove any remaining leading/trailing whitespace
content = content.strip()
return content
def save_final_results(self, state: Dict):
"""Save all final results."""
output_dir = state["output_dir"]
# Save all analysis results and diagram codes
for key, value in state.items():
# Skip metadata fields
if key in ["project_name", "project_description", "output_dir", "selected_diagrams"]:
continue
if value: # Only save if there's content
self.save_step_result(output_dir, key, value)
# Save metadata
metadata = {
"project_name": state["project_name"],
"project_description": state["project_description"],
"generated_diagrams": state["selected_diagrams"],
"timestamp": datetime.now().isoformat()
}
import json
with open(os.path.join(output_dir, "metadata.json"), "w") as f:
json.dump(metadata, f, indent=2) |