Add file upload + landing page

app.py

@@ -2,8 +2,8 @@
 SecureHeal Agent — HuggingFace Space FastAPI Server
 ────────────────────────────────────────────────────
 Loads the trained model at startup, caches it, and exposes a FastAPI
-endpoint that takes application code → runs the SecureHeal agent →
-finds vulnerabilities → suggests fixes → returns structured response.
+endpoint that takes application code (text or file upload) → runs the
+SecureHeal agent → finds vulnerabilities → suggests fixes.
 
 Deploy to HF Spaces with GPU (T4).
 """
@@ -13,26 +13,24 @@ import json
 import re
 import torch
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
 from typing import Optional, List
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import pipeline as hf_pipeline
 
 # ────────────────────── Model Cache ──────────────────────────
 
 MODEL_ID = os.environ.get("MODEL_ID", "Nitesh-Reddy/secureheal-agent-v2")
-PIPE = None  # Global pipeline — loaded once at startup
+PIPE = None
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """Load model at startup, keep in memory for all requests."""
     global PIPE
     print(f"🔄 Loading model: {MODEL_ID}")
-    print("   This takes ~2 min on first load, then cached...")
-
-    PIPE = pipeline(
+    PIPE = hf_pipeline(
         "text-generation",
         model=MODEL_ID,
         torch_dtype=torch.float16,
@@ -40,7 +38,6 @@ async def lifespan(app: FastAPI):
     )
     print(f"✅ Model loaded and cached!")
     yield
-    print("🛑 Shutting down...")
 
 
 # ────────────────────── FastAPI App ──────────────────────────
@@ -60,7 +57,7 @@ app.add_middleware(
 )
 
 
-# ────────────────────── Request/Response Models ──────────────
+# ────────────────────── Models ───────────────────────────────
 
 class ScanRequest(BaseModel):
     code: str
@@ -75,7 +72,7 @@ class VulnerabilityReport(BaseModel):
     vulnerabilities_found: bool
     tool_calls: List[ToolCall]
     analysis: str
-    raw_output: str
+    filename: Optional[str] = None
 
 class AgentRequest(BaseModel):
     prompt: str
@@ -86,22 +83,18 @@ class AgentResponse(BaseModel):
     tool_calls: List[ToolCall]
 
 
-# ────────────────────── Helper: Parse Tool Calls ─────────────
+# ────────────────────── Helper ───────────────────────────────
 
 def parse_tool_calls(text: str) -> List[ToolCall]:
-    """Extract <tool_call>tool_name({...})</tool_call> from model output."""
     calls = []
     pattern = r'<tool_call>\s*(\w+)\((\{.*?\})\)\s*</tool_call>'
     matches = re.findall(pattern, text, re.DOTALL)
-
     for tool_name, args_str in matches:
         try:
             args = json.loads(args_str)
         except json.JSONDecodeError:
             args = {"raw": args_str}
         calls.append(ToolCall(tool=tool_name, args=args))
-
-    # Fallback: find tool mentions without proper wrapping
     if not calls:
         valid_tools = [
             "scan_code", "simulate_attack", "apply_patch", "run_tests",
@@ -110,83 +103,147 @@ def parse_tool_calls(text: str) -> List[ToolCall]:
         for tool in valid_tools:
             if tool in text.lower():
                 calls.append(ToolCall(tool=tool, args={}))
-
     return calls
 
 
+def run_agent(code: str, context: str = "web application", max_tokens: int = 512) -> str:
+    """Run the SecureHeal agent on the given code."""
+    prompt = (
+        f"You are an autonomous SRE and Security agent. "
+        f"Analyze the following {context} code for vulnerabilities. "
+        f"Use scan_code, simulate_attack, apply_patch, run_tests to analyze and fix. "
+        f'Output each action as <tool_call>tool_name({{"param": "value"}})</tool_call>. '
+        f"End with DONE when finished.\n\n"
+        f"Code to analyze:\n```\n{code}\n```"
+    )
+    messages = [{"role": "user", "content": prompt}]
+    output = PIPE(messages, max_new_tokens=max_tokens, do_sample=True, temperature=0.7)
+    return output[0]["generated_text"][-1]["content"]
+
+
 # ────────────────────── Endpoints ────────────────────────────
 
-@app.get("/")
+@app.get("/", response_class=HTMLResponse)
 async def root():
-    return {
-        "service": "SecureHeal Agent",
-        "model": MODEL_ID,
-        "status": "ready" if PIPE else "loading",
-        "endpoints": {
-            "/scan": "POST — Scan code for vulnerabilities",
-            "/agent": "POST — Free-form agent prompt",
-            "/health": "GET — Health check",
-        },
-    }
+    """Landing page with usage instructions."""
+    return """
+    <html>
+    <head><title>SecureHeal Agent</title>
+    <style>
+        body { font-family: system-ui; max-width: 800px; margin: 40px auto; padding: 20px;
+               background: #0d1117; color: #e6edf3; }
+        h1 { color: #58a6ff; }
+        code { background: #161b22; padding: 2px 6px; border-radius: 4px; color: #f0883e; }
+        pre { background: #161b22; padding: 16px; border-radius: 8px; overflow-x: auto; }
+        .endpoint { background: #161b22; padding: 12px 16px; border-radius: 8px;
+                     border-left: 3px solid #58a6ff; margin: 12px 0; }
+        a { color: #58a6ff; }
+    </style></head>
+    <body>
+        <h1>🛡️ SecureHeal Agent API</h1>
+        <p>Autonomous SRE & Security agent — trained with GRPO on Llama 3 8B</p>
+
+        <h2>Endpoints</h2>
+
+        <div class="endpoint">
+            <strong>POST /scan</strong> — Scan code (JSON body)<br>
+            <code>{"code": "your code here", "context": "web app"}</code>
+        </div>
+
+        <div class="endpoint">
+            <strong>POST /scan/file</strong> — Upload a file to scan<br>
+            <code>curl -F "file=@app.py" -F "context=flask app" URL/scan/file</code>
+        </div>
+
+        <div class="endpoint">
+            <strong>POST /agent</strong> — Free-form agent prompt<br>
+            <code>{"prompt": "Find SQL injection in login function"}</code>
+        </div>
+
+        <div class="endpoint">
+            <strong>GET /health</strong> — Health check
+        </div>
+
+        <h2>Example</h2>
+        <pre>curl -X POST /scan/file \\
+  -F "file=@vulnerable_app.py" \\
+  -F "context=flask web application"</pre>
+
+        <p>Model: <a href="https://huggingface.co/Nitesh-Reddy/secureheal-agent-v2">Nitesh-Reddy/secureheal-agent-v2</a></p>
+        <p><a href="/docs">📖 Interactive API Docs (Swagger)</a></p>
+    </body></html>
+    """
 
 
 @app.get("/health")
 async def health():
-    return {"status": "healthy", "model_loaded": PIPE is not None}
+    return {"status": "healthy", "model_loaded": PIPE is not None, "model": MODEL_ID}
 
 
 @app.post("/scan", response_model=VulnerabilityReport)
-async def scan_code(request: ScanRequest):
+async def scan_code_json(request: ScanRequest):
+    """Scan code for vulnerabilities (JSON body with code string)."""
+    if not PIPE:
+        raise HTTPException(503, "Model still loading")
+
+    response_text = run_agent(request.code, request.context, request.max_tokens)
+    tool_calls = parse_tool_calls(response_text)
+
+    return VulnerabilityReport(
+        vulnerabilities_found=len(tool_calls) > 0,
+        tool_calls=tool_calls,
+        analysis=response_text,
+    )
+
+
+@app.post("/scan/file", response_model=VulnerabilityReport)
+async def scan_code_file(
+    file: UploadFile = File(..., description="Source code file to scan"),
+    context: str = Form("web application", description="What kind of app (e.g. flask, django, express)"),
+    max_tokens: int = Form(512, description="Max response tokens"),
+):
     """
-    Scan application code for vulnerabilities.
-    The agent analyzes the code and returns structured tool calls + fixes.
+    Upload a source code file for vulnerability scanning.
+    Supports .py, .js, .ts, .java, .go, .rb, .php, etc.
     """
     if not PIPE:
-        raise HTTPException(503, "Model still loading, try again in ~2 min")
+        raise HTTPException(503, "Model still loading")
 
-    prompt = (
-        f"You are an autonomous SRE and Security agent. "
-        f"Analyze the following {request.context} code for vulnerabilities. "
-        f"Use scan_code, simulate_attack, apply_patch, run_tests to analyze and fix. "
-        f"Output each action as <tool_call>tool_name({{\"param\": \"value\"}})</tool_call>. "
-        f"End with DONE when finished.\n\n"
-        f"Code to analyze:\n```\n{request.code}\n```"
-    )
+    # Read file content
+    content = await file.read()
+    try:
+        code = content.decode("utf-8")
+    except UnicodeDecodeError:
+        raise HTTPException(400, "File must be a text/source code file")
 
-    messages = [{"role": "user", "content": prompt}]
-    output = PIPE(messages, max_new_tokens=request.max_tokens, do_sample=True, temperature=0.7)
-    response_text = output[0]["generated_text"][-1]["content"]
+    # Truncate very long files to fit model context
+    if len(code) > 8000:
+        code = code[:8000] + "\n\n# ... (truncated, file too large)"
 
+    response_text = run_agent(code, context, max_tokens)
     tool_calls = parse_tool_calls(response_text)
 
     return VulnerabilityReport(
         vulnerabilities_found=len(tool_calls) > 0,
         tool_calls=tool_calls,
         analysis=response_text,
-        raw_output=response_text,
+        filename=file.filename,
     )
 
 
 @app.post("/agent", response_model=AgentResponse)
 async def agent_prompt(request: AgentRequest):
-    """
-    Send a free-form prompt to the SecureHeal agent.
-    """
+    """Send a free-form prompt to the SecureHeal agent."""
     if not PIPE:
-        raise HTTPException(503, "Model still loading, try again in ~2 min")
+        raise HTTPException(503, "Model still loading")
 
     messages = [{"role": "user", "content": request.prompt}]
     output = PIPE(messages, max_new_tokens=request.max_tokens, do_sample=True, temperature=0.7)
     response_text = output[0]["generated_text"][-1]["content"]
     tool_calls = parse_tool_calls(response_text)
 
-    return AgentResponse(
-        response=response_text,
-        tool_calls=tool_calls,
-    )
-
+    return AgentResponse(response=response_text, tool_calls=tool_calls)
 
-# ────────────────────── Run ──────────────────────────────────
 
 if __name__ == "__main__":
     import uvicorn

requirements.txt

@@ -5,3 +5,4 @@ bitsandbytes
 fastapi
 uvicorn[standard]
 pydantic>=2.0
+python-multipart