|
|
import json |
|
|
import os |
|
|
from typing import List |
|
|
from fastapi import FastAPI, HTTPException |
|
|
from fastapi.responses import HTMLResponse |
|
|
from pydantic import BaseModel |
|
|
from model import generate_structure |
|
|
|
|
|
app = FastAPI() |
|
|
STRUCTURE_FILE = "structures.json" |
|
|
structure_store = [] |
|
|
|
|
|
class Prompt(BaseModel): |
|
|
text: str |
|
|
|
|
|
class Structure(BaseModel): |
|
|
name: str |
|
|
description: str |
|
|
required_blocks: List[str] |
|
|
tags: List[str] |
|
|
|
|
|
@app.post("/generate") |
|
|
def generate(prompt: Prompt): |
|
|
result = generate_structure(prompt.text) |
|
|
return {"response": result} |
|
|
|
|
|
@app.post("/register_structure") |
|
|
def register_structure(structure: Structure): |
|
|
|
|
|
structures = [] |
|
|
if os.path.exists(STRUCTURE_FILE): |
|
|
with open(STRUCTURE_FILE, "r") as f: |
|
|
try: |
|
|
structures = json.load(f) |
|
|
except json.JSONDecodeError: |
|
|
pass |
|
|
|
|
|
|
|
|
for s in structures: |
|
|
if s["name"] == structure.name: |
|
|
raise HTTPException(status_code=400, detail="Structure with this name already exists") |
|
|
|
|
|
|
|
|
structures.append(structure.dict()) |
|
|
with open(STRUCTURE_FILE, "w") as f: |
|
|
json.dump(structures, f, indent=2) |
|
|
|
|
|
return {"message": "Structure registered successfully"} |
|
|
|
|
|
@app.get("/structures", response_class=HTMLResponse) |
|
|
def list_structures(): |
|
|
html = "<h2>Registered Structures</h2><table border='1'><tr><th>Name</th><th>Description</th><th>Required Blocks</th><th>Tags</th></tr>" |
|
|
for structure in structure_store: |
|
|
html += f"<tr><td>{structure.name}</td><td>{structure.description}</td><td>{', '.join(structure.required_blocks)}</td><td>{', '.join(structure.tags)}</td></tr>" |
|
|
html += "</table>" |
|
|
return html |