app/codegen/gradio_generator.py

"""
Generate Gradio app code using the HF Inference API (LLM-powered).
Falls back to template-based generation if the API is unavailable.
"""

import os
import re
from huggingface_hub import InferenceClient


LLM_MODEL = "Qwen/Qwen2.5-72B-Instruct"

# Pre-built templates for common app types (used as fallback and as LLM context)
GRADIO_TEMPLATES = {
    "chatbot": '''import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("{model_id}")

def respond(message, history, system_message, max_tokens, temperature, top_p):
    messages = [{{"role": "system", "content": system_message}}]
    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({{"role": "user", "content": user_msg}})
        if bot_msg:
            messages.append({{"role": "assistant", "content": bot_msg}})
    messages.append({{"role": "user", "content": message}})

    response = ""
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = chunk.choices[0].delta.content or ""
        response += token
        yield response

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a helpful assistant.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max tokens"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    title="{title}",
    description="{description}",
)

if __name__ == "__main__":
    demo.launch()
''',

    "image_classifier": '''import gradio as gr
from huggingface_hub import InferenceClient
import json

client = InferenceClient("{model_id}")

def classify_image(image):
    if image is None:
        return {{}}
    result = client.image_classification(image)
    return {{item["label"]: round(item["score"], 4) for item in result}}

demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="filepath", label="Upload Image"),
    outputs=gr.Label(num_top_classes=5, label="Predictions"),
    title="{title}",
    description="{description}",
    examples=[],
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()
''',

    "text_summarizer": '''import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("{model_id}")

def summarize(text, max_length, min_length):
    if not text.strip():
        return "Please enter some text to summarize."
    result = client.summarization(
        text,
        parameters={{"max_length": int(max_length), "min_length": int(min_length)}},
    )
    return result["summary_text"] if isinstance(result, dict) else result[0]["summary_text"]

demo = gr.Interface(
    fn=summarize,
    inputs=[
        gr.Textbox(lines=10, placeholder="Paste your text here...", label="Input Text"),
        gr.Slider(50, 500, value=150, step=10, label="Max Summary Length"),
        gr.Slider(10, 100, value=30, step=5, label="Min Summary Length"),
    ],
    outputs=gr.Textbox(lines=5, label="Summary"),
    title="{title}",
    description="{description}",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()
''',

    "sentiment_analyzer": '''import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("{model_id}")

def analyze_sentiment(text):
    if not text.strip():
        return {{}}, ""
    result = client.text_classification(text)
    scores = {{item["label"]: round(item["score"], 4) for item in result}}
    top_label = max(scores, key=scores.get)
    explanation = f"The text is predominantly **{{top_label}}** (confidence: {{scores[top_label]:.1%}})"
    return scores, explanation

demo = gr.Interface(
    fn=analyze_sentiment,
    inputs=gr.Textbox(lines=5, placeholder="Enter text to analyze...", label="Input Text"),
    outputs=[
        gr.Label(label="Sentiment Scores"),
        gr.Markdown(label="Analysis"),
    ],
    title="{title}",
    description="{description}",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()
''',

    "text_generator": '''import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("{model_id}")

def generate_text(prompt, max_tokens, temperature, top_p):
    if not prompt.strip():
        return "Please enter a prompt."
    messages = [{{"role": "user", "content": prompt}}]
    response = ""
    for chunk in client.chat_completion(
        messages, max_tokens=int(max_tokens), stream=True,
        temperature=temperature, top_p=top_p,
    ):
        token = chunk.choices[0].delta.content or ""
        response += token
        yield response

demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=3, placeholder="Enter your prompt...", label="Prompt"),
        gr.Slider(50, 2048, value=512, step=50, label="Max Tokens"),
        gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    outputs=gr.Textbox(lines=10, label="Generated Text"),
    title="{title}",
    description="{description}",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()
''',

    "translator": '''import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("{model_id}")

def translate(text, target_language):
    if not text.strip():
        return "Please enter text to translate."
    prompt = f"Translate the following text to {{target_language}}:\\n\\n{{text}}"
    messages = [{{"role": "user", "content": prompt}}]
    result = client.chat_completion(messages, max_tokens=1024)
    return result.choices[0].message.content

LANGUAGES = ["French", "Spanish", "German", "Italian", "Portuguese",
             "Chinese", "Japanese", "Korean", "Arabic", "Hindi", "Russian"]

demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.Textbox(lines=5, placeholder="Enter text to translate...", label="Input Text"),
        gr.Dropdown(choices=LANGUAGES, value="French", label="Target Language"),
    ],
    outputs=gr.Textbox(lines=5, label="Translation"),
    title="{title}",
    description="{description}",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()
''',

    "question_answering": '''import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("{model_id}")

def answer_question(context, question):
    if not context.strip() or not question.strip():
        return "Please provide both context and a question."
    result = client.question_answering(question=question, context=context)
    answer = result["answer"]
    score = result["score"]
    return f"**Answer:** {{answer}}\\n\\n**Confidence:** {{score:.1%}}"

demo = gr.Interface(
    fn=answer_question,
    inputs=[
        gr.Textbox(lines=8, placeholder="Paste the context here...", label="Context"),
        gr.Textbox(lines=2, placeholder="Ask a question...", label="Question"),
    ],
    outputs=gr.Markdown(label="Answer"),
    title="{title}",
    description="{description}",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()
''',
}


class GradioGenerator:
    """Generate Gradio app.py code for a given plan."""

    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) -> str:
        """Generate app.py content for a Gradio Space."""
        template_key = plan.get("template_key")
        model_id = self._get_model_id(plan)
        title = plan.get("title", "My App")
        description = plan.get("description", "")

        # First try LLM generation for best results
        try:
            code = self._generate_with_llm(plan, prompt, model_id, title, description)
            if code and len(code) > 100:
                return code
        except Exception:
            pass

        # Fallback to template-based generation
        if template_key and template_key in GRADIO_TEMPLATES:
            return GRADIO_TEMPLATES[template_key].format(
                model_id=model_id,
                title=title,
                description=description,
            )

        # Final fallback: generic Gradio app
        return self._generate_generic(plan, model_id, title, description)

    def _generate_with_llm(
        self, plan: dict, prompt: str, model_id: str, title: str, description: str
    ) -> str:
        """Use the LLM to generate custom Gradio app code."""
        template_key = plan.get("template_key")
        reference_code = ""
        if template_key and template_key in GRADIO_TEMPLATES:
            reference_code = GRADIO_TEMPLATES[template_key].format(
                model_id=model_id, title=title, description=description,
            )

        system_prompt = """You are an expert Python developer specializing in Gradio applications for Hugging Face Spaces.
You generate complete, working app.py files that are production-ready.

Rules:
1. Output ONLY the Python code, no explanations or markdown.
2. Use `huggingface_hub.InferenceClient` for model inference (NOT transformers or pipeline).
3. Always include proper error handling.
4. The app must use `demo.launch()` at the end.
5. Include descriptive title and description in the interface.
6. Use appropriate Gradio components for the use case.
7. Code must be complete and runnable as-is.
8. Do NOT use `transformers`, `torch`, or `tensorflow` imports - use only `huggingface_hub.InferenceClient`.
9. For streaming text generation, use `client.chat_completion(..., stream=True)`.
10. For image tasks, use `client.image_classification()`, `client.object_detection()`, etc.
11. For text tasks, use `client.text_classification()`, `client.summarization()`, etc."""

        user_prompt = f"""Generate a complete Gradio app.py for this request:

USER REQUEST: {prompt}

APP PLAN:
- SDK: gradio
- Type: {plan.get('app_type', 'custom')}
- Model: {model_id}
- Model Task: {plan.get('model_task', 'text-generation')}
- Components: {', '.join(plan.get('components', []))}
- Title: {title}
- Description: {description}
- Extra Features: {', '.join(plan.get('extra_features', []))}
"""

        if reference_code:
            user_prompt += f"""
REFERENCE TEMPLATE (adapt and improve this based on the user's specific request):
```python
{reference_code}
```
"""

        user_prompt += "\nGenerate the complete app.py code now:"

        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._extract_code(raw)

    def _extract_code(self, text: str) -> str:
        """Extract Python code from LLM response."""
        # Try to find code blocks
        code_blocks = re.findall(r"```(?:python)?\s*\n(.*?)```", text, re.DOTALL)
        if code_blocks:
            # Return the longest code block
            return max(code_blocks, key=len).strip()

        # If no code blocks, check if the whole response looks like Python code
        lines = text.strip().split("\n")
        if lines and (lines[0].startswith("import ") or lines[0].startswith("from ")):
            return text.strip()

        return text.strip()

    def _get_model_id(self, plan: dict) -> str:
        """Get the best model ID from the plan."""
        models = plan.get("recommended_models", [])
        if models:
            return models[0]["id"]

        # Fallback defaults by task
        task_defaults = {
            "text-generation": "Qwen/Qwen2.5-7B-Instruct",
            "text-classification": "cardiffnlp/twitter-roberta-base-sentiment-latest",
            "summarization": "facebook/bart-large-cnn",
            "translation": "Qwen/Qwen2.5-7B-Instruct",
            "image-classification": "google/vit-base-patch16-224",
            "object-detection": "facebook/detr-resnet-50",
            "text-to-image": "stabilityai/stable-diffusion-xl-base-1.0",
            "automatic-speech-recognition": "openai/whisper-base",
            "question-answering": "deepset/roberta-base-squad2",
        }
        return task_defaults.get(plan.get("model_task", ""), "Qwen/Qwen2.5-7B-Instruct")

    def _generate_generic(self, plan: dict, model_id: str, title: str, description: str) -> str:
        """Generate a generic Gradio app when no template matches."""
        task = plan.get("model_task", "text-generation")

        if task in ("text-generation",):
            return GRADIO_TEMPLATES["text_generator"].format(
                model_id=model_id, title=title, description=description,
            )
        elif task in ("text-classification",):
            return GRADIO_TEMPLATES["sentiment_analyzer"].format(
                model_id=model_id, title=title, description=description,
            )
        elif task in ("summarization",):
            return GRADIO_TEMPLATES["text_summarizer"].format(
                model_id=model_id, title=title, description=description,
            )
        elif task in ("image-classification",):
            return GRADIO_TEMPLATES["image_classifier"].format(
                model_id=model_id, title=title, description=description,
            )
        elif task in ("question-answering",):
            return GRADIO_TEMPLATES["question_answering"].format(
                model_id=model_id, title=title, description=description,
            )
        else:
            return GRADIO_TEMPLATES["text_generator"].format(
                model_id=model_id, title=title, description=description,
            )

    def edit(self, plan: dict, current_code: str, edit_prompt: str) -> str:
        """Edit existing Gradio code based on user instructions."""
        try:
            system_prompt = """You are an expert Python developer. You will receive existing Gradio app code and an edit request.
Modify the code according to the request and return the COMPLETE updated code.
Output ONLY the Python code, no explanations or markdown fences.
Keep all existing functionality unless the user explicitly asks to remove something.
Use huggingface_hub.InferenceClient for model inference."""

            user_prompt = f"""EXISTING CODE:
```python
{current_code}
```

EDIT REQUEST: {edit_prompt}

Return the complete updated app.py code:"""

            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 = self._extract_code(raw)
            if code and len(code) > 50:
                return code
        except Exception:
            pass

        return current_code