Deploy SecureHeal Agent API

Dockerfile

@@ -0,0 +1,12 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app.py .
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,12 +1,9 @@
 ---
-title: Secureheal Agent
-emoji: 🏢
+title: SecureHeal Agent
+emoji: 🛡️
 colorFrom: red
 colorTo: purple
-sdk: gradio
-sdk_version: 6.13.0
-app_file: app.py
-pinned: false
+sdk: docker
+app_port: 7860
+pinned: true
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,193 @@
+"""
+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.
+
+Deploy to HF Spaces with GPU (T4).
+"""
+
+import os
+import json
+import re
+import torch
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Optional, List
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+# ────────────────────── Model Cache ──────────────────────────
+
+MODEL_ID = os.environ.get("MODEL_ID", "Nitesh-Reddy/secureheal-agent-v2")
+PIPE = None  # Global pipeline — loaded once at startup
+
+
+@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(
+        "text-generation",
+        model=MODEL_ID,
+        torch_dtype=torch.float16,
+        device_map="auto",
+    )
+    print(f"✅ Model loaded and cached!")
+    yield
+    print("🛑 Shutting down...")
+
+
+# ────────────────────── FastAPI App ──────────────────────────
+
+app = FastAPI(
+    title="SecureHeal Agent API",
+    description="Autonomous SRE & Security agent — scans code, finds vulnerabilities, suggests fixes",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+# ────────────────────── Request/Response Models ──────────────
+
+class ScanRequest(BaseModel):
+    code: str
+    context: Optional[str] = "web application"
+    max_tokens: Optional[int] = 512
+
+class ToolCall(BaseModel):
+    tool: str
+    args: dict
+
+class VulnerabilityReport(BaseModel):
+    vulnerabilities_found: bool
+    tool_calls: List[ToolCall]
+    analysis: str
+    raw_output: str
+
+class AgentRequest(BaseModel):
+    prompt: str
+    max_tokens: Optional[int] = 512
+
+class AgentResponse(BaseModel):
+    response: str
+    tool_calls: List[ToolCall]
+
+
+# ────────────────────── Helper: Parse Tool Calls ─────────────
+
+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",
+            "restart_service", "clean_data", "reallocate_resources", "classify_issue",
+        ]
+        for tool in valid_tools:
+            if tool in text.lower():
+                calls.append(ToolCall(tool=tool, args={}))
+
+    return calls
+
+
+# ────────────────────── Endpoints ────────────────────────────
+
+@app.get("/")
+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",
+        },
+    }
+
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy", "model_loaded": PIPE is not None}
+
+
+@app.post("/scan", response_model=VulnerabilityReport)
+async def scan_code(request: ScanRequest):
+    """
+    Scan application code for vulnerabilities.
+    The agent analyzes the code and returns structured tool calls + fixes.
+    """
+    if not PIPE:
+        raise HTTPException(503, "Model still loading, try again in ~2 min")
+
+    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```"
+    )
+
+    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"]
+
+    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,
+    )
+
+
+@app.post("/agent", response_model=AgentResponse)
+async def agent_prompt(request: AgentRequest):
+    """
+    Send a free-form prompt to the SecureHeal agent.
+    """
+    if not PIPE:
+        raise HTTPException(503, "Model still loading, try again in ~2 min")
+
+    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,
+    )
+
+
+# ────────────────────── Run ──────────────────────────────────
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -0,0 +1,7 @@
+transformers>=4.46
+torch>=2.0
+accelerate
+bitsandbytes
+fastapi
+uvicorn[standard]
+pydantic>=2.0