Update app.py

app.py

@@ -1,161 +1,212 @@
 # app.py
-# Simple coding assistant with analyze-and-rewrite workflow
-# Brand: odyssey
-
+# Odyssey Coder — enhanced: multi-file output + analyze & rewrite + download
 import os
+import tempfile
 import requests
 import gradio as gr
+from typing import Tuple
 
-# ===== DEFAULTS =====
-DEFAULT_MODEL = "stabilityai/stable-code-instruct-3b"
-DEFAULT_MAX_NEW_TOKENS = 512
-DEFAULT_TEMPERATURE = 0.2
-DEFAULT_TOP_P = 0.95
-DEFAULT_TOP_K = 50
-
+# -------------------
+# Configuration
+# -------------------
+DEFAULT_MODEL = "bigcode/starcoder2-3b"  # پیش‌فرض سبک‌تر
 HF_INFERENCE_URL = "https://api-inference.huggingface.co/models/{}"
-
-# ===== Helper: call Hugging Face text-generation inference API =====
-def hf_text_generation(model: str, prompt: str, hf_token: str, params: dict):
+HF_TOKEN = os.environ.get("HF_API_TOKEN")  # از Secrets خوانده می‌شود
+
+# Preset model choices (user can override with custom model field)
+PRESET_MODELS = [
+    "bigcode/starcoder2-3b",
+    "Salesforce/codegen-2B-multi",
+    "huggingface/CodeParrot-small",
+    "google/flan-t5-small"  # fallback NLP model (not code-specialized)
+]
+
+# Allowed file extensions for download
+ALLOWED_EXTS = [".py", ".js", ".html", ".md", ".txt"]
+
+# -------------------
+# Helper: call HF Inference API
+# -------------------
+def hf_text_generation(model: str, prompt: str, hf_token: str, params: dict) -> Tuple[bool, str]:
+    """
+    Returns (ok, text_or_error)
+    ok=True => text result
+    ok=False => error message
+    """
     url = HF_INFERENCE_URL.format(model)
-    headers = {"Authorization": f"Bearer {hf_token}"}
+    headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
     payload = {"inputs": prompt, "parameters": params}
-    r = requests.post(url, headers=headers, json=payload, timeout=120)
-    r.raise_for_status()
-    data = r.json()
-
-    if isinstance(data, dict) and "error" in data:
-        raise RuntimeError(f"Hugging Face error: {data['error']}")
 
-    # Common response shapes:
-    # 1) [{"generated_text": "..."}]
-    # 2) [{"some_key": ...}] or string-like results
-    if isinstance(data, list) and len(data) > 0:
-        first = data[0]
-        if isinstance(first, dict) and "generated_text" in first:
-            return first["generated_text"]
-        return str(first)
-
-    return str(data)
-
-# ===== Analysis & rewrite instruction prompt =====
-ANALYSIS_INSTRUCTION = """
-You are a senior code reviewer and refactorer.
-Step 1: Analyze the following input (it can be source code or a prompt). Provide concise bullet points: bugs, edge-cases, security concerns, performance issues, missing pieces, and concrete suggestions.
-Step 2: Produce an optimized, rewritten version of the code or a clarified prompt that is ready-to-run or ready-to-feed into a code generation model.
-
-Return the output in two clearly marked sections:
+    try:
+        r = requests.post(url, headers=headers, json=payload, timeout=120)
+        r.raise_for_status()
+        data = r.json()
+        # typical shapes: [{"generated_text": "..."}]
+        if isinstance(data, list) and len(data) > 0:
+            first = data[0]
+            if isinstance(first, dict) and "generated_text" in first:
+                return True, first["generated_text"]
+            # sometimes models return plain text or dict with other keys
+            return True, str(first)
+        return True, str(data)
+    except requests.exceptions.HTTPError as he:
+        status = he.response.status_code if he.response is not None else None
+        if status == 401:
+            return False, "Authentication error: invalid or missing HF_API_TOKEN."
+        if status == 403:
+            return False, "Permission error: the selected model or resource is restricted."
+        if status == 404:
+            return False, "Model not found (404). Check the model id."
+        if status == 410:
+            return False, "Model endpoint not available (410 Gone). Try another model."
+        # generic
+        return False, f"HTTP error {status}: {he.response.text if he.response is not None else str(he)}"
+    except requests.exceptions.RequestException as e:
+        return False, f"Request error: {str(e)}"
+
+# -------------------
+# Core: analyze & rewrite
+# -------------------
+ANALYSIS_INSTRUCTION = """You are an expert senior code reviewer and refactorer.
+Step 1: Provide concise bullet-point analysis: bugs, edge-cases, security issues, missing pieces, and suggestions.
+Step 2: Produce an optimized, runnable, and well-commented version of the code or a clarified prompt.
+
+Return ONLY two sections marked exactly as below:
 
 ===ANALYSIS===
-- bullet points...
+- ...
 
 ===REWRITTEN===
-<optimized code or clarified prompt>
-
-Return only these two sections and nothing else.
+<rewritten code or prompt>
 """
 
-def analyze_and_rewrite_via_api(model, user_input, hf_token, analysis_max_tokens=512, gen_max_tokens=512, temperature=0.15):
+def analyze_and_rewrite(model: str, user_input: str, hf_token: str, max_new_tokens: int, temperature: float, top_p: float, top_k: int):
     prompt = ANALYSIS_INSTRUCTION + "\n\nINPUT:\n" + user_input + "\n\n"
     params = {
-        "max_new_tokens": analysis_max_tokens,
+        "max_new_tokens": max_new_tokens,
         "temperature": temperature,
-        "top_p": DEFAULT_TOP_P,
-        "top_k": DEFAULT_TOP_K,
+        "top_p": top_p,
+        "top_k": top_k,
         "return_full_text": False
     }
-    raw = hf_text_generation(model, prompt, hf_token, params)
+    ok, resp = hf_text_generation(model, prompt, hf_token, params)
+    return ok, resp
 
-    if "===ANALYSIS===" in raw and "===REWRITTEN===" in raw:
-        analysis = raw.split("===ANALYSIS===")[1].split("===REWRITTEN===")[0].strip()
-        rewritten = raw.split("===REWRITTEN===")[1].strip()
-        return analysis, rewritten, raw
-
-    # Fallback: return full raw as analysis if markers are missing
-    return raw, "", raw
-
-def generate_code_via_api(model, prompt, hf_token, max_new_tokens=256, temperature=0.2, top_p=0.95, top_k=50):
+def normal_generate(model: str, user_input: str, hf_token: str, max_new_tokens: int, temperature: float, top_p: float, top_k: int):
     params = {
-        "max_new_tokens": int(max_new_tokens),
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "top_k": int(top_k),
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
         "return_full_text": False
     }
-    return hf_text_generation(model, prompt, hf_token, params)
-
-# ===== Gradio UI =====
-with gr.Blocks(title="odyssey — Coding Assistant & Rewriter") as demo:
-    gr.Markdown("## odyssey — Coding Assistant and Rewriter (Qwen Coder)")
-    gr.Markdown(
-        "This tool lets you customize model, token, temperature and output length. "
-        "Enable 'Analyze → Rewrite' to have the model first analyze the input and then produce an optimized rewrite."
-    )
-
+    return hf_text_generation(model, user_input, hf_token, params)
+
+# -------------------
+# Utility: write output to a file and return path
+# -------------------
+def write_temp_file(content: str, filename: str, ext: str) -> str:
+    if not ext.startswith("."):
+        ext = "." + ext
+    # safe filename
+    safe_name = "".join(c for c in filename if c.isalnum() or c in ("-", "_", "." )).strip()
+    if not safe_name:
+        safe_name = "output"
+    if not safe_name.endswith(ext):
+        safe_name = safe_name + ext
+    tmpdir = tempfile.mkdtemp()
+    path = os.path.join(tmpdir, safe_name)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(content)
+    return path
+
+# -------------------
+# Gradio app logic
+# -------------------
+def on_generate(inp, mode_analyze, preset_model, custom_model, use_preset, hf_token_input,
+                max_new_tokens, temperature, top_p, top_k, file_ext, filename, make_download):
+    # Resolve model
+    model = preset_model if use_preset else (custom_model.strip() or DEFAULT_MODEL)
+    hf_token = hf_token_input or os.environ.get("HF_API_TOKEN")
+    if not hf_token:
+        return ("ERROR: No HF token set. Add HF_API_TOKEN to Space Secrets or enter token here.", "", "", None)
+
+    if mode_analyze:
+        ok, resp = analyze_and_rewrite(model, inp, hf_token, max_new_tokens, temperature, top_p, top_k)
+        if not ok:
+            return (f"Error during analysis: {resp}", "", "", None)
+        # parse sections
+        if "===ANALYSIS===" in resp and "===REWRITTEN===" in resp:
+            analysis = resp.split("===ANALYSIS===")[1].split("===REWRITTEN===")[0].strip()
+            rewritten = resp.split("===REWRITTEN===")[1].strip()
+        else:
+            # fallback: show full response in analysis
+            analysis = resp
+            rewritten = ""
+        # If rewritten exists, generate final output from rewritten (normal generate)
+        final_prompt = rewritten if rewritten.strip() else inp
+        ok2, gen = normal_generate(model, final_prompt, hf_token, max_new_tokens, temperature, top_p, top_k)
+        if not ok2:
+            return (f"Error during generation: {gen}", analysis, rewritten, None)
+    else:
+        # normal generation
+        ok, gen = normal_generate(model, inp, hf_token, max_new_tokens, temperature, top_p, top_k)
+        if not ok:
+            return (f"Error during generation: {gen}", "", "", None)
+        analysis = ""
+        rewritten = ""
+        gen = gen
+
+    # If user asked for download, create file
+    file_path = None
+    if make_download:
+        ext = file_ext if file_ext.startswith(".") else f".{file_ext}"
+        if ext not in ALLOWED_EXTS:
+            return ("Invalid file extension selected.", analysis, rewritten, None)
+        file_path = write_temp_file(gen, filename or "odyssey_output", ext)
+
+    return ("OK", analysis, rewritten, file_path)
+
+# -------------------
+# Build Gradio UI
+# -------------------
+with gr.Blocks(title="odyssey — Enhanced Coder (multi-file)") as demo:
+    gr.Markdown("## odyssey — Enhanced Coding Assistant\nAnalyze → Rewrite mode, choose output file type and download the generated file.")
     with gr.Row():
         with gr.Column(scale=2):
-            user_input = gr.Textbox(
-                label="Input (code or prompt)",
-                lines=8,
-                placeholder="Paste the code you want reviewed/refactored or the prompt you want optimized..."
-            )
-            thinking_toggle = gr.Checkbox(label="Analyze → Rewrite (thinking mode)", value=True)
-            model_input = gr.Textbox(label="Hugging Face model name", value=DEFAULT_MODEL)
-            use_remote = gr.Checkbox(label="Use Hugging Face Inference API (recommended)", value=True)
-            hf_token_input = gr.Textbox(
-                label="Hugging Face API token (leave empty to use HF_API_TOKEN secret)",
-                type="password"
-            )
-            max_tokens_input = gr.Slider(
-                minimum=64, maximum=2048, step=16,
-                label="max_new_tokens", value=DEFAULT_MAX_NEW_TOKENS
-            )
-            temp_input = gr.Slider(
-                minimum=0.0, maximum=1.0, step=0.01,
-                label="temperature", value=DEFAULT_TEMPERATURE
-            )
-            top_p_input = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, label="top_p", value=DEFAULT_TOP_P)
-            top_k_input = gr.Slider(minimum=0, maximum=200, step=1, label="top_k", value=DEFAULT_TOP_K)
-            generate_btn = gr.Button("Run")
-
+            inp = gr.Textbox(label="Input (code or prompt)", lines=8, placeholder="Paste code or prompt...")
+            mode_analyze = gr.Checkbox(label="Analyze → Rewrite (recommended)", value=True)
+            # model selection
+            preset_model = gr.Dropdown(choices=PRESET_MODELS, value=PRESET_MODELS[0], label="Preset model")
+            use_preset = gr.Checkbox(label="Use preset model (otherwise use custom model below)", value=True)
+            custom_model = gr.Textbox(label="Custom model id (optional)", placeholder="e.g. username/model-name")
+            hf_token_input = gr.Textbox(label="HuggingFace API Token (leave empty to use HF_API_TOKEN secret)", type="password")
+
+            with gr.Row():
+                max_new_tokens = gr.Slider(64, 2048, value=512, step=64, label="max_new_tokens")
+                temperature = gr.Slider(0.0, 1.0, value=0.2, step=0.01, label="temperature")
+            with gr.Row():
+                top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="top_p")
+                top_k = gr.Slider(0, 200, value=50, step=1, label="top_k")
+
+            gr.Markdown("### Output file options")
+            file_ext = gr.Dropdown(choices=[e.lstrip(".") for e in ALLOWED_EXTS], value="py", label="File extension")
+            filename = gr.Textbox(label="Filename (without extension)", value="odyssey_output")
+            make_download = gr.Checkbox(label="Create downloadable file", value=True)
+
+            btn = gr.Button("Generate")
         with gr.Column(scale=2):
-            analysis_out = gr.Textbox(label="Model Analysis", lines=8)
-            rewritten_out = gr.Textbox(label="Rewritten Code / Prompt", lines=12)
-            normal_out = gr.Textbox(label="Normal Generation Output", lines=12)
-
-    def on_generate(inp, thinking, model_name, use_remote_flag, token, max_new_toks, temp, top_p, top_k):
-        hf_token = token or os.environ.get("HF_API_TOKEN")
-        if use_remote_flag and not hf_token:
-            return "ERROR: Please provide a Hugging Face API token or set HF_API_TOKEN in Secrets.", "", ""
-
-        try:
-            if thinking:
-                analysis, rewritten, raw = analyze_and_rewrite_via_api(
-                    model_name, inp, hf_token,
-                    analysis_max_tokens=min(1024, max_new_toks // 2),
-                    gen_max_tokens=max_new_toks,
-                    temperature=temp
-                )
-                final_source = rewritten if rewritten.strip() else inp
-                normal_generation = generate_code_via_api(
-                    model_name, final_source, hf_token,
-                    max_new_tokens=max_new_toks,
-                    temperature=temp,
-                    top_p=top_p, top_k=top_k
-                )
-                return analysis, rewritten, normal_generation
-            else:
-                out = generate_code_via_api(model_name, inp, hf_token,
-                                           max_new_tokens=max_new_toks, temperature=temp,
-                                           top_p=top_p, top_k=top_k)
-                return "", "", out
-        except Exception as e:
-            return f"Model execution error: {str(e)}", "", ""
-
-    generate_btn.click(
+            status_out = gr.Textbox(label="Status / Errors", lines=2)
+            analysis_out = gr.Textbox(label="Analysis (if analyze mode)", lines=8)
+            rewritten_out = gr.Textbox(label="Rewritten code / prompt", lines=12)
+            gen_out = gr.Textbox(label="Final generated text", lines=14)
+            file_down = gr.File(label="Download generated file")
+
+    btn.click(
         on_generate,
-        inputs=[user_input, thinking_toggle, model_input, use_remote, hf_token_input, max_tokens_input, temp_input, top_p_input, top_k_input],
-        outputs=[analysis_out, rewritten_out, normal_out]
+        inputs=[inp, mode_analyze, preset_model, custom_model, use_preset, hf_token_input,
+                max_new_tokens, temperature, top_p, top_k, file_ext, filename, make_download],
+        outputs=[status_out, analysis_out, rewritten_out, file_down]
     )
 
 if __name__ == "__main__":