Spaces:

ruslanmv
/

autoapp-builder

Sleeping

File size: 11,372 Bytes

2c304fc

"""
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

    @property
    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