Spaces:
Sleeping
Sleeping
| """ | |
| Generate Docker-based (FastAPI) Space code using LLM with template fallbacks. | |
| """ | |
| import os | |
| import re | |
| from huggingface_hub import InferenceClient | |
| LLM_MODEL = "Qwen/Qwen2.5-72B-Instruct" | |
| DOCKER_TEMPLATES = { | |
| "rest_api": { | |
| "app.py": '''from fastapi import FastAPI, HTTPException | |
| from fastapi.responses import HTMLResponse | |
| from pydantic import BaseModel | |
| from huggingface_hub import InferenceClient | |
| import os | |
| app = FastAPI( | |
| title="{title}", | |
| description="{description}", | |
| version="1.0.0", | |
| ) | |
| client = InferenceClient("{model_id}", token=os.environ.get("HF_TOKEN")) | |
| class TextRequest(BaseModel): | |
| text: str | |
| max_tokens: int = 512 | |
| temperature: float = 0.7 | |
| class TextResponse(BaseModel): | |
| result: str | |
| model: str | |
| class SummarizeRequest(BaseModel): | |
| text: str | |
| max_length: int = 150 | |
| min_length: int = 30 | |
| @app.get("/", response_class=HTMLResponse) | |
| async def root(): | |
| return """ | |
| <html> | |
| <head><title>{title}</title></head> | |
| <body style="font-family: system-ui; max-width: 800px; margin: 50px auto; padding: 20px;"> | |
| <h1>{title}</h1> | |
| <p>{description}</p> | |
| <h2>Endpoints</h2> | |
| <ul> | |
| <li><code>POST /generate</code> - Generate text</li> | |
| <li><code>POST /summarize</code> - Summarize text</li> | |
| <li><code>GET /health</code> - Health check</li> | |
| <li><code>GET /docs</code> - API documentation</li> | |
| </ul> | |
| </body> | |
| </html> | |
| """ | |
| @app.post("/generate", response_model=TextResponse) | |
| async def generate_text(request: TextRequest): | |
| try: | |
| messages = [{{"role": "user", "content": request.text}}] | |
| response = client.chat_completion( | |
| messages, | |
| max_tokens=request.max_tokens, | |
| temperature=request.temperature, | |
| ) | |
| return TextResponse( | |
| result=response.choices[0].message.content, | |
| model="{model_id}", | |
| ) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| @app.post("/summarize") | |
| async def summarize_text(request: SummarizeRequest): | |
| try: | |
| prompt = f"Summarize the following text in a concise way:\\n\\n{{request.text}}" | |
| messages = [{{"role": "user", "content": prompt}}] | |
| response = client.chat_completion(messages, max_tokens=request.max_length) | |
| return {{"summary": response.choices[0].message.content, "model": "{model_id}"}} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| @app.get("/health") | |
| async def health_check(): | |
| return {{"status": "healthy", "model": "{model_id}"}} | |
| ''', | |
| "Dockerfile": '''FROM python:3.11-slim | |
| WORKDIR /app | |
| COPY requirements.txt . | |
| RUN pip install --no-cache-dir -r requirements.txt | |
| COPY . . | |
| RUN useradd -m -u 1000 appuser | |
| USER appuser | |
| EXPOSE 7860 | |
| CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"] | |
| ''', | |
| "requirements.txt": "fastapi==0.115.6\nuvicorn[standard]==0.34.0\nhuggingface-hub==0.27.1\n", | |
| }, | |
| "generic_docker": { | |
| "app.py": '''from fastapi import FastAPI, HTTPException | |
| from fastapi.responses import HTMLResponse | |
| from pydantic import BaseModel | |
| from huggingface_hub import InferenceClient | |
| import os | |
| app = FastAPI(title="{title}", version="1.0.0") | |
| client = InferenceClient("{model_id}", token=os.environ.get("HF_TOKEN")) | |
| class QueryRequest(BaseModel): | |
| query: str | |
| max_tokens: int = 512 | |
| @app.get("/", response_class=HTMLResponse) | |
| async def root(): | |
| return """ | |
| <html> | |
| <head><title>{title}</title> | |
| <style> | |
| body {{ font-family: system-ui; max-width: 900px; margin: 50px auto; padding: 20px; background: #0f172a; color: #e2e8f0; }} | |
| h1 {{ color: #38bdf8; }} | |
| code {{ background: #1e293b; padding: 2px 8px; border-radius: 4px; }} | |
| a {{ color: #38bdf8; }} | |
| </style> | |
| </head> | |
| <body> | |
| <h1>{title}</h1> | |
| <p>{description}</p> | |
| <h2>API Documentation</h2> | |
| <p>Visit <a href="/docs">/docs</a> for interactive API documentation.</p> | |
| <h2>Endpoints</h2> | |
| <ul> | |
| <li><code>POST /query</code> - Process a query</li> | |
| <li><code>GET /health</code> - Health check</li> | |
| </ul> | |
| </body> | |
| </html> | |
| """ | |
| @app.post("/query") | |
| async def process_query(request: QueryRequest): | |
| try: | |
| messages = [{{"role": "user", "content": request.query}}] | |
| response = client.chat_completion(messages, max_tokens=request.max_tokens) | |
| return {{ | |
| "response": response.choices[0].message.content, | |
| "model": "{model_id}", | |
| }} | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=str(e)) | |
| @app.get("/health") | |
| async def health(): | |
| return {{"status": "healthy"}} | |
| ''', | |
| "Dockerfile": '''FROM python:3.11-slim | |
| WORKDIR /app | |
| COPY requirements.txt . | |
| RUN pip install --no-cache-dir -r requirements.txt | |
| COPY . . | |
| RUN useradd -m -u 1000 appuser | |
| USER appuser | |
| EXPOSE 7860 | |
| CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"] | |
| ''', | |
| "requirements.txt": "fastapi==0.115.6\nuvicorn[standard]==0.34.0\nhuggingface-hub==0.27.1\n", | |
| }, | |
| } | |
| class DockerGenerator: | |
| """Generate Docker-based Space files.""" | |
| def __init__(self): | |
| self._client = None | |
| def client(self) -> InferenceClient: | |
| if self._client is None: | |
| token = os.environ.get("HF_TOKEN", None) | |
| self._client = InferenceClient(LLM_MODEL, token=token) | |
| return self._client | |
| def generate(self, plan: dict, prompt: str) -> dict: | |
| """ | |
| Generate all files for a Docker-based Space. | |
| Returns dict of {filename: content}. | |
| """ | |
| model_id = self._get_model_id(plan) | |
| title = plan.get("title", "My API") | |
| description = plan.get("description", "") | |
| template_key = plan.get("template_key", "generic_docker") | |
| # Try LLM generation first | |
| try: | |
| files = self._generate_with_llm(plan, prompt, model_id, title, description) | |
| if files and "app.py" in files and len(files["app.py"]) > 100: | |
| # Ensure required files exist | |
| if "Dockerfile" not in files: | |
| files["Dockerfile"] = DOCKER_TEMPLATES["generic_docker"]["Dockerfile"] | |
| if "requirements.txt" not in files: | |
| files["requirements.txt"] = DOCKER_TEMPLATES["generic_docker"]["requirements.txt"] | |
| return files | |
| except Exception: | |
| pass | |
| # Fallback to templates | |
| template = DOCKER_TEMPLATES.get(template_key, DOCKER_TEMPLATES["generic_docker"]) | |
| files = {} | |
| for filename, content in template.items(): | |
| files[filename] = content.format( | |
| model_id=model_id, | |
| title=title, | |
| description=description, | |
| ) | |
| return files | |
| def _generate_with_llm( | |
| self, plan: dict, prompt: str, model_id: str, title: str, description: str | |
| ) -> dict: | |
| """Use LLM to generate custom Docker app files.""" | |
| system_prompt = """You are an expert Python developer specializing in FastAPI applications for Hugging Face Docker Spaces. | |
| Generate complete, production-ready code files. | |
| Rules: | |
| 1. Output files in this exact format - each file delimited by markers: | |
| === FILENAME: app.py === | |
| (file content) | |
| === FILENAME: requirements.txt === | |
| (file content) | |
| === FILENAME: Dockerfile === | |
| (file content) | |
| 2. Use FastAPI for the web framework. | |
| 3. Use `huggingface_hub.InferenceClient` for model inference. | |
| 4. Include a root GET endpoint that returns an HTML page describing the API. | |
| 5. Include a /health endpoint. | |
| 6. Include proper error handling. | |
| 7. The Dockerfile should use python:3.11-slim, create a non-root user, and expose port 7860. | |
| 8. Use uvicorn to serve the app on 0.0.0.0:7860. | |
| 9. Do NOT use transformers, torch, or tensorflow.""" | |
| user_prompt = f"""Generate a complete Docker-based HF Space for this request: | |
| USER REQUEST: {prompt} | |
| APP PLAN: | |
| - Type: {plan.get('app_type', 'custom')} | |
| - Model: {model_id} | |
| - Model Task: {plan.get('model_task', 'text-generation')} | |
| - Title: {title} | |
| - Description: {description} | |
| - Extra Features: {', '.join(plan.get('extra_features', []))} | |
| Generate the files (app.py, requirements.txt, Dockerfile):""" | |
| response = self.client.chat_completion( | |
| messages=[ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| max_tokens=4096, | |
| temperature=0.3, | |
| ) | |
| raw = response.choices[0].message.content | |
| return self._parse_files(raw) | |
| def _parse_files(self, text: str) -> dict: | |
| """Parse LLM output into a dict of filename -> content.""" | |
| files = {} | |
| # Try marker-based parsing first | |
| parts = re.split(r"===\s*FILENAME:\s*(.+?)\s*===", text) | |
| if len(parts) >= 3: | |
| for i in range(1, len(parts), 2): | |
| filename = parts[i].strip() | |
| content = parts[i + 1].strip() if i + 1 < len(parts) else "" | |
| # Remove markdown code fences if present | |
| content = re.sub(r"^```\w*\n?", "", content) | |
| content = re.sub(r"\n?```$", "", content) | |
| files[filename] = content.strip() + "\n" | |
| # Try code block parsing as fallback | |
| if not files: | |
| blocks = re.findall( | |
| r"(?:#|//)\s*(\S+\.(?:py|txt|dockerfile))\s*\n```\w*\n(.*?)```", | |
| text, re.DOTALL | re.IGNORECASE, | |
| ) | |
| for name, content in blocks: | |
| files[name.strip()] = content.strip() + "\n" | |
| return files | |
| def _get_model_id(self, plan: dict) -> str: | |
| models = plan.get("recommended_models", []) | |
| if models: | |
| return models[0]["id"] | |
| return "Qwen/Qwen2.5-7B-Instruct" | |
| def edit(self, plan: dict, current_files: dict, edit_prompt: str) -> dict: | |
| """Edit existing Docker app files.""" | |
| try: | |
| app_code = current_files.get("app.py", "") | |
| system_prompt = """You are an expert Python developer. Modify the existing FastAPI app code according to the edit request. | |
| Output ONLY the updated Python code for app.py, no explanations or markdown fences. | |
| Keep all existing functionality unless explicitly asked to remove something. | |
| Use huggingface_hub.InferenceClient for model inference.""" | |
| user_prompt = f"""EXISTING app.py: | |
| ```python | |
| {app_code} | |
| ``` | |
| EDIT REQUEST: {edit_prompt} | |
| Return the complete updated app.py:""" | |
| response = self.client.chat_completion( | |
| messages=[ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ], | |
| max_tokens=4096, | |
| temperature=0.2, | |
| ) | |
| raw = response.choices[0].message.content | |
| code = re.sub(r"^```\w*\n?", "", raw.strip()) | |
| code = re.sub(r"\n?```$", "", code) | |
| if code and len(code) > 50: | |
| updated = dict(current_files) | |
| updated["app.py"] = code.strip() + "\n" | |
| return updated | |
| except Exception: | |
| pass | |
| return current_files | |