image_generators.py
1 """Image-generator registry CRUD endpoints. 2 3 Mirrors `restai/routers/llms.py` so the admin UX (list / create / edit / 4 delete + team grants) carries over without surprises. The registry holds 5 both: 6 7 - **Local** workers (auto-seeded on startup from `restai/image/workers/*`). 8 Always selectable; admin can flip `enabled` and rename for display, but 9 cannot delete (re-seeded next boot). 10 - **External** providers — `openai` (incl. OpenAI-spec compatibles via 11 `options.base_url`) and `google` (Imagen / Nano Banana). Created freely 12 by the admin with per-row encrypted credentials. 13 14 API key fields in `options` are masked as `"********"` on read; the PATCH 15 handler preserves the existing value when it sees that sentinel back. 16 """ 17 import json 18 import logging 19 import traceback 20 from typing import Optional 21 22 from fastapi import APIRouter, Depends, HTTPException, Path, Request 23 24 from restai import config 25 from restai.auth import get_current_username, get_current_username_admin 26 from restai.database import DBWrapper, get_db_wrapper 27 from restai.models.databasemodels import ImageGeneratorDatabase 28 from restai.models.models import ( 29 ImageGeneratorModel, 30 ImageGeneratorModelCreate, 31 ImageGeneratorModelUpdate, 32 User, 33 ) 34 35 logging.basicConfig(level=config.LOG_LEVEL) 36 37 router = APIRouter() 38 39 40 _SENSITIVE_OPT_KEYS = {"api_key", "key", "password", "secret"} 41 42 43 def _mask_options(options: Optional[dict]) -> Optional[dict]: 44 """Replace sensitive fields with the `"********"` sentinel before 45 serializing to the client. Same set the encrypt helpers use.""" 46 if not options: 47 return options 48 try: 49 masked = dict(options) 50 for k in _SENSITIVE_OPT_KEYS: 51 if k in masked and masked[k]: 52 masked[k] = "********" 53 return masked 54 except Exception: 55 return options 56 57 58 @router.get("/image_generators", response_model=list[ImageGeneratorModel]) 59 async def list_image_generators( 60 user: User = Depends(get_current_username), 61 db_wrapper: DBWrapper = Depends(get_db_wrapper), 62 ): 63 """List image generators. Non-admins see only those granted to a team 64 they're a member of (matches the LLM listing pattern).""" 65 rows = db_wrapper.get_image_generators() 66 67 if not user.is_admin: 68 allowed_names = set() 69 for team in user.teams or []: 70 for ig in (team.image_generators or []): 71 allowed_names.add(getattr(ig, "generator_name", ig)) 72 rows = [r for r in rows if r.name in allowed_names] 73 74 out: list[ImageGeneratorModel] = [] 75 for r in rows: 76 m = ImageGeneratorModel.model_validate(r) 77 m.options = _mask_options(m.options) 78 out.append(m) 79 return out 80 81 82 @router.get("/image_generators/{generator_id}", response_model=ImageGeneratorModel) 83 async def get_image_generator( 84 generator_id: int = Path(description="Image generator ID"), 85 _: User = Depends(get_current_username), 86 db_wrapper: DBWrapper = Depends(get_db_wrapper), 87 ): 88 row = db_wrapper.get_image_generator_by_id(generator_id) 89 if row is None: 90 raise HTTPException(status_code=404, detail="Image generator not found") 91 m = ImageGeneratorModel.model_validate(row) 92 m.options = _mask_options(m.options) 93 return m 94 95 96 @router.post("/image_generators", status_code=201, response_model=ImageGeneratorModel) 97 async def create_image_generator( 98 body: ImageGeneratorModelCreate, 99 _: User = Depends(get_current_username_admin), 100 db_wrapper: DBWrapper = Depends(get_db_wrapper), 101 ): 102 """Register a new image generator (admin only).""" 103 if db_wrapper.get_image_generator_by_name(body.name): 104 raise HTTPException(status_code=409, detail=f"Image generator '{body.name}' already exists") 105 if body.class_name == "local": 106 # Local generators are auto-seeded; admin shouldn't create them 107 # by hand (the name has to match a real worker module). 108 raise HTTPException( 109 status_code=400, 110 detail="Local generators are auto-discovered from restai/image/workers/*; you can't create them manually.", 111 ) 112 try: 113 opts = body.options if isinstance(body.options, dict) else (json.loads(body.options) if body.options else {}) 114 row = db_wrapper.create_image_generator( 115 name=body.name, 116 class_name=body.class_name, 117 options=opts, 118 privacy=body.privacy, 119 description=body.description, 120 enabled=body.enabled, 121 ) 122 m = ImageGeneratorModel.model_validate(row) 123 m.options = _mask_options(m.options) 124 return m 125 except HTTPException: 126 raise 127 except Exception as e: 128 logging.error(e) 129 traceback.print_tb(e.__traceback__) 130 raise HTTPException(status_code=500, detail=f"Failed to create image generator '{body.name}'") 131 132 133 @router.patch("/image_generators/{generator_id}", response_model=ImageGeneratorModel) 134 async def update_image_generator( 135 request: Request, 136 generator_id: int = Path(description="Image generator ID"), 137 body: ImageGeneratorModelUpdate = ..., 138 _: User = Depends(get_current_username_admin), 139 db_wrapper: DBWrapper = Depends(get_db_wrapper), 140 ): 141 """Update an image generator (admin only). For local generators the 142 `enabled`/`description`/`privacy` fields are accepted; `class_name` and 143 `options` changes are ignored — those come from the worker module.""" 144 row: Optional[ImageGeneratorDatabase] = db_wrapper.get_image_generator_by_id(generator_id) 145 if row is None: 146 raise HTTPException(status_code=404, detail="Image generator not found") 147 148 if row.class_name == "local": 149 # Strip class_name + options changes for local rows so an admin 150 # can't accidentally point a local row at an external provider. 151 body.class_name = None 152 body.options = None 153 154 db_wrapper.edit_image_generator(row, body) 155 m = ImageGeneratorModel.model_validate(row) 156 m.options = _mask_options(m.options) 157 return m 158 159 160 @router.delete("/image_generators/{generator_id}") 161 async def delete_image_generator( 162 generator_id: int = Path(description="Image generator ID"), 163 _: User = Depends(get_current_username_admin), 164 db_wrapper: DBWrapper = Depends(get_db_wrapper), 165 ): 166 """Delete an image generator (admin only). Local generators cannot be 167 deleted — they would just be re-seeded on the next boot. Disable them 168 via `enabled=false` instead.""" 169 row: Optional[ImageGeneratorDatabase] = db_wrapper.get_image_generator_by_id(generator_id) 170 if row is None: 171 raise HTTPException(status_code=404, detail="Image generator not found") 172 if row.class_name == "local": 173 raise HTTPException( 174 status_code=400, 175 detail=f"Cannot delete local generator '{row.name}'. Set enabled=false instead.", 176 ) 177 name = row.name 178 db_wrapper.delete_image_generator(row) 179 return {"deleted": name}