Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import uuid | |
| import os | |
| import json | |
| import shutil | |
| import zipfile | |
| from typing import Dict, List | |
| from spec_validator import validate_and_clean_spec | |
| from pipeline import run_pipeline | |
| # ===================================================== | |
| # CONFIG | |
| # ===================================================== | |
| BASE_RUN_DIR = "runs" | |
| SCHEMA_PATH = "spec_schema.json" | |
| st.set_page_config( | |
| page_title="AI REST API Generator", | |
| layout="wide", | |
| ) | |
| # ===================================================== | |
| # HELPERS | |
| # ===================================================== | |
| def create_run_dir(): | |
| run_id = uuid.uuid4().hex[:8] | |
| run_dir = os.path.join(BASE_RUN_DIR, f"run_{run_id}") | |
| os.makedirs(run_dir, exist_ok=True) | |
| return run_id, run_dir | |
| def zip_folder(folder_path: str) -> str: | |
| zip_path = f"{folder_path}.zip" | |
| with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: | |
| for root, _, files in os.walk(folder_path): | |
| for file in files: | |
| full_path = os.path.join(root, file) | |
| arcname = os.path.relpath(full_path, folder_path) | |
| zipf.write(full_path, arcname) | |
| return zip_path | |
| def load_schema(): | |
| with open(SCHEMA_PATH, "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| # ===================================================== | |
| # SESSION STATE INIT | |
| # ===================================================== | |
| if "spec" not in st.session_state: | |
| st.session_state.spec = { | |
| "project_name": "", | |
| "description": "", | |
| "database": {}, | |
| "auth": {"type": "jwt"}, | |
| "api_config": {"base_url": "/api/"}, | |
| "apps": {"core": {"models": {}, "apis": {}}}, | |
| "use_ai_models": True, | |
| "llm": {} | |
| } | |
| if "models_ui" not in st.session_state: | |
| st.session_state.models_ui = [] | |
| if "logs" not in st.session_state: | |
| st.session_state.logs = [] | |
| # ===================================================== | |
| # UI HEADER | |
| # ===================================================== | |
| st.title("🤖 AI REST API Generator") | |
| st.markdown( | |
| """ | |
| **Build Django REST APIs using AI or manual configuration.** | |
| 🚧 *This website is under development.* | |
| """ | |
| ) | |
| st.divider() | |
| # ===================================================== | |
| # SECTION 1 — LLM CONFIGURATION | |
| # ===================================================== | |
| st.header("🔑 LLM Configuration") | |
| llm_provider = st.selectbox( | |
| "Select LLM Provider", | |
| ["Groq", "OpenAI"] | |
| ) | |
| api_key = st.text_input( | |
| f"{llm_provider} API Key", | |
| type="password", | |
| help="Your API key is used only for this session and never stored." | |
| ) | |
| st.session_state.spec["llm"] = { | |
| "provider": llm_provider.lower() | |
| } | |
| # ===================================================== | |
| # SECTION 2 — PROJECT BASICS | |
| # ===================================================== | |
| st.header("📦 Project Information") | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| st.session_state.spec["project_name"] = st.text_input( | |
| "Project Name", | |
| placeholder="my_backend_project" | |
| ) | |
| with col2: | |
| st.session_state.spec["description"] = st.text_area( | |
| "Project Description", | |
| placeholder="Short description of your backend" | |
| ) | |
| # ===================================================== | |
| # SECTION 3 — DATABASE CONFIGURATION | |
| # ===================================================== | |
| st.header("🗄 Database Configuration") | |
| db_engine = st.selectbox( | |
| "Database Engine", | |
| ["sqlite", "postgresql", "mysql"] | |
| ) | |
| if db_engine == "sqlite": | |
| st.session_state.spec["database"] = { | |
| "engine": "sqlite", | |
| "name": "db.sqlite3" | |
| } | |
| else: | |
| col1, col2 = st.columns(2) | |
| with col1: | |
| db_name = st.text_input("Database Name") | |
| db_user = st.text_input("Database User") | |
| db_password = st.text_input("Database Password", type="password") | |
| with col2: | |
| db_host = st.text_input("Database Host", value="localhost") | |
| db_port = st.number_input("Database Port", value=5432) | |
| st.session_state.spec["database"] = { | |
| "engine": db_engine, | |
| "name": db_name, | |
| "user": db_user, | |
| "password": db_password, | |
| "host": db_host, | |
| "port": db_port | |
| } | |
| # ===================================================== | |
| # SECTION 4 — MODEL MODE | |
| # ===================================================== | |
| st.header("🧱 Model Generation") | |
| use_ai = st.radio( | |
| "How do you want to create models?", | |
| ["AI Generated Models", "Manual Model Builder"] | |
| ) | |
| st.session_state.spec["use_ai_models"] = (use_ai == "AI Generated Models") | |
| # ===================================================== | |
| # SECTION 5 — MANUAL MODEL BUILDER | |
| # ===================================================== | |
| if not st.session_state.spec["use_ai_models"]: | |
| st.subheader("🛠 Manual Model Builder") | |
| if st.button("➕ Add Model"): | |
| st.session_state.models_ui.append({ | |
| "name": "", | |
| "fields": [] | |
| }) | |
| model_names = [] | |
| for mi, model in enumerate(st.session_state.models_ui): | |
| with st.expander(f"Model {mi + 1}", expanded=True): | |
| model["name"] = st.text_input( | |
| "Model Name", | |
| model["name"], | |
| key=f"model_{mi}" | |
| ) | |
| if model["name"]: | |
| model_names.append(model["name"]) | |
| if st.button("➕ Add Column", key=f"add_col_{mi}"): | |
| model["fields"].append({ | |
| "name": "", | |
| "type": "CharField", | |
| "primary_key": False, | |
| "unique": False, | |
| "null": False, | |
| "relation": None, | |
| "on_delete": "CASCADE" | |
| }) | |
| for fi, field in enumerate(model["fields"]): | |
| cols = st.columns(7) | |
| field["name"] = cols[0].text_input( | |
| "Column", | |
| field["name"], | |
| key=f"fname_{mi}_{fi}" | |
| ) | |
| field["type"] = cols[1].selectbox( | |
| "Type", | |
| [ | |
| "CharField", | |
| "IntegerField", | |
| "UUIDField", | |
| "BooleanField", | |
| "DateField", | |
| "OneToOne", | |
| "OneToMany" | |
| ], | |
| key=f"ftype_{mi}_{fi}" | |
| ) | |
| field["primary_key"] = cols[2].checkbox( | |
| "PK", | |
| key=f"fpk_{mi}_{fi}" | |
| ) | |
| field["unique"] = cols[3].checkbox( | |
| "Unique", | |
| key=f"funq_{mi}_{fi}" | |
| ) | |
| field["null"] = cols[4].checkbox( | |
| "Null", | |
| key=f"fnull_{mi}_{fi}" | |
| ) | |
| if field["type"] in ["OneToOne", "OneToMany"]: | |
| field["relation"] = cols[5].selectbox( | |
| "Reference Model", | |
| model_names, | |
| key=f"frel_{mi}_{fi}" | |
| ) | |
| field["on_delete"] = cols[6].selectbox( | |
| "On Delete", | |
| ["CASCADE", "SET_NULL", "PROTECT"], | |
| key=f"fdel_{mi}_{fi}" | |
| ) | |
| # Convert UI → spec | |
| models_spec = {} | |
| apis_spec = {} | |
| for model in st.session_state.models_ui: | |
| if not model["name"]: | |
| continue | |
| fields_spec = {} | |
| for f in model["fields"]: | |
| if not f["name"]: | |
| continue | |
| if f["type"] == "OneToOne": | |
| fields_spec[f["name"]] = { | |
| "type": "OneToOneField", | |
| "to": f["relation"], | |
| "on_delete": f["on_delete"], | |
| "null": f["null"], | |
| "unique": True | |
| } | |
| elif f["type"] == "OneToMany": | |
| fields_spec[f["name"]] = { | |
| "type": "ForeignKey", | |
| "to": f["relation"], | |
| "on_delete": f["on_delete"], | |
| "null": f["null"] | |
| } | |
| else: | |
| fields_spec[f["name"]] = { | |
| "type": f["type"], | |
| "primary_key": f["primary_key"], | |
| "unique": f["unique"], | |
| "null": f["null"] | |
| } | |
| models_spec[model["name"]] = {"fields": fields_spec} | |
| apis_spec[model["name"]] = ["list", "create", "retrieve", "update", "delete"] | |
| st.session_state.spec["apps"]["core"]["models"] = models_spec | |
| st.session_state.spec["apps"]["core"]["apis"] = apis_spec | |
| # ===================================================== | |
| # SECTION 6 — GENERATION | |
| # ===================================================== | |
| st.header("🚀 Generate Project") | |
| # UI elements (define once, above button) | |
| progress_bar = st.progress(0) | |
| log_box = st.empty() | |
| def log(msg: str): | |
| st.session_state.logs.append(msg) | |
| log_box.code( | |
| "\n".join(st.session_state.logs), | |
| language="text" | |
| ) | |
| def ui_log(msg): | |
| st.session_state.logs.append(msg) | |
| log_box.write("\n".join(st.session_state.logs)) | |
| def ui_progress(value): | |
| progress_bar.progress(value) | |
| if st.button("Generate Backend"): | |
| # ---------------------------- | |
| # 1️⃣ Basic validation | |
| # ---------------------------- | |
| if not api_key: | |
| st.error("API key is required to generate the project.") | |
| st.stop() | |
| st.session_state.logs.clear() | |
| progress_bar.progress(0) | |
| # ---------------------------- | |
| # 2️⃣ Load schema & validate UI spec | |
| # ---------------------------- | |
| schema = load_schema() | |
| valid, cleaned_spec, errors, warnings = validate_and_clean_spec( | |
| st.session_state.spec, | |
| schema | |
| ) | |
| if not valid: | |
| st.error("Specification validation failed.") | |
| st.json(errors) | |
| st.stop() | |
| # ---------------------------- | |
| # 3️⃣ Create isolated run directory | |
| # ---------------------------- | |
| run_id, run_dir = create_run_dir() | |
| ui_log(f"🆔 Run ID: {run_id}") | |
| ui_log("🚀 Starting generation pipeline...") | |
| # ---------------------------- | |
| # 4️⃣ Run backend pipeline | |
| # ---------------------------- | |
| try: | |
| run_pipeline( | |
| spec=cleaned_spec, | |
| run_dir=run_dir, | |
| llm_provider=llm_provider.lower(), | |
| api_key=api_key, | |
| log_callback=ui_log, | |
| progress_callback=ui_progress | |
| ) | |
| ui_log("📦 Zipping project...") | |
| zip_path = zip_folder(run_dir) | |
| ui_log("✅ Project ready for download") | |
| with open(zip_path, "rb") as f: | |
| st.download_button( | |
| "⬇ Download Generated Project", | |
| f, | |
| file_name=f"{cleaned_spec['project_name']}.zip" | |
| ) | |
| # ---------------------------- | |
| # 5️⃣ Error handling | |
| # ---------------------------- | |
| except Exception as e: | |
| ui_log("❌ Pipeline failed") | |
| ui_log(str(e)) | |
| st.error(f"Generation failed at step: {e}") | |
| # ---------------------------- | |
| # 6️⃣ Cleanup (safe) | |
| # ---------------------------- | |
| finally: | |
| shutil.rmtree(run_dir, ignore_errors=True) | |
| # ===================================================== | |
| # FOOTER | |
| # ===================================================== | |
| st.divider() | |
| st.markdown( | |
| """ | |
| 📧 **Contact:** harshadhole04@gmail.com | |
| © AI REST API Generator — Under Development | |
| """ | |
| ) | |