griddava/pull-request-validator

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from fastapi.responses import HTMLResponse4import requests5import os6 7 8# ----------------------------9# 1. Configuration10# ----------------------------11 12# Remove hardcoded API key and use an environment variable13HF_API_KEY = os.getenv("HF_API_KEY")14 15if not HF_API_KEY:16    raise RuntimeError("Hugging Face API key is not set. Please set the HF_API_KEY environment variable.")17 18HF_MODEL_NAME = "gpt2"  # A reliable text generation model available on HF Inference API19 20# ----------------------------21# 2. FastAPI App Initialization22# ----------------------------23 24app = FastAPI(25    title="AI Code Review Service",26    description="An API to get AI-powered code reviews for pull request diffs.",27    version="1.0.0",28)29 30# ----------------------------31# 3. No Local Model Loading (Using HF API Instead)32# ----------------------------33 34@app.on_event("startup")35async def startup_event():36    """37    On server startup, validate HF API key.38    """39    print("Server starting up...")40    print(f"Using Hugging Face API with model: {HF_MODEL_NAME}")41    if not HF_API_KEY:42        print("WARNING: HF_API_KEY not set!")43    else:44        print("HF_API_KEY is configured.")45 46# ----------------------------47# 4. API Request/Response Models48# ----------------------------49 50class ReviewRequest(BaseModel):51    diff: str52 53class ReviewComment(BaseModel):54    file_path: str55    line_number: int56    comment_text: str57 58class ReviewResponse(BaseModel):59    comments: list[ReviewComment]60 61# ----------------------------62# 5. The AI Review Logic63# ----------------------------64 65def run_ai_inference(diff: str) -> str:66    """67    Sends the code diff to Hugging Face Inference API to get the review.68    """69    # Better prompt for meaningful completions70    prompt = f"""Code review feedback:71 72{diff[:200]}73 74Feedback: This code could be improved by"""75 76    headers = {77        "Authorization": f"Bearer {HF_API_KEY}",78        "Content-Type": "application/json"79    }80    payload = {81        "inputs": prompt,82        "parameters": {83            "max_new_tokens": 32,84            "temperature": 0.7,85            "top_p": 0.986        }87    }88 89    try:90        response = requests.post(91            f"https://api-inference.huggingface.co/models/{HF_MODEL_NAME}",92            headers=headers,93            json=payload,94            timeout=3095        )96 97        if response.status_code != 200:98            print(f"HF API Error: {response.status_code} - {response.text}")99            return "Consider adding proper documentation and error handling."100 101        response_data = response.json()102        print(f"HF API Response: {response_data}")103        104        if isinstance(response_data, list) and len(response_data) > 0:105            generated_text = response_data[0].get("generated_text", "")106            # Extract only the new generated part (after our prompt)107            if generated_text.startswith(prompt):108                response_text = generated_text[len(prompt):].strip()109            else:110                response_text = generated_text.strip()111        else:112            response_text = "Unable to generate a meaningful review."113 114    except Exception as e:115        print(f"HF API Exception: {e}")116        return "Consider adding proper documentation and error handling."117 118    # Clean up the response119    response_text = response_text.strip()120 121    # Handle different completion patterns122    if response_text.startswith("error handling"):123        review = "Consider adding error handling and input validation."124    elif response_text.startswith("documentation"):125        review = "Consider adding documentation and type hints."126    elif response_text.startswith("input validation"):127        review = "Consider adding input validation and error checks."128    elif response_text.startswith("type hints"):129        review = "Consider adding type hints and documentation."130    else:131        # Extract meaningful content132        lines = [line.strip() for line in response_text.split('\n') if line.strip()]133        if lines and len(lines[0]) > 3:134            first_line = lines[0]135            # Clean up common artifacts136            if first_line.startswith('#'):137                first_line = first_line[1:].strip()138            if len(first_line) > 10:139                review = f"Consider adding {first_line.lower()}."140            else:141                review = "Consider adding proper documentation and error handling."142        else:143            review = "Consider adding proper documentation and error handling."144 145    return review146 147def parse_ai_response(response_text: str) -> list[ReviewComment]:148    """149    Parses the raw text from the AI to extract the JSON array.150    """151    # For codegen-350M-mono, just wrap the review in a single comment152    return [ReviewComment(153        file_path="code_reviewed.py",154        line_number=1,155        comment_text=response_text.strip()156    )]157 158# ----------------------------159# 6. The API Endpoint160# ----------------------------161 162@app.post("/review", response_model=ReviewResponse)163async def get_code_review(request: ReviewRequest):164    if not request.diff:165        raise HTTPException(status_code=400, detail="Diff content cannot be empty.")166 167    import time168    start_time = time.time()169    print(f"Starting review request at {start_time}")170 171    try:172        print("Running AI inference...")173        ai_response_text = run_ai_inference(request.diff)174        print(f"AI inference completed in {time.time() - start_time:.2f} seconds")175        176        print("Parsing AI response...")177        parsed_comments = parse_ai_response(ai_response_text)178        print(f"Total processing time: {time.time() - start_time:.2f} seconds")179        180        return ReviewResponse(comments=parsed_comments)181 182    except Exception as e:183        print(f"An unexpected error occurred after {time.time() - start_time:.2f} seconds: {e}")184        raise HTTPException(status_code=500, detail="An internal error occurred while processing the review.")185 186# ----------------------------187# 7. Health Check Endpoint188# ----------------------------189@app.get("/", response_class=HTMLResponse)190def root_html():191    """Return HTML for browser viewing."""192    return """193    <!DOCTYPE html>194    <html>195    <head>196        <title>AI Code Review Service</title>197        <style>198            body { font-family: Arial, sans-serif; margin: 40px; }199            .status { color: green; font-weight: bold; }200            .endpoint { background: #f4f4f4; padding: 10px; margin: 10px 0; border-radius: 5px; }201        </style>202    </head>203    <body>204        <h1>AI Code Review Service</h1>205        <p class="status">✅ Service is running with AI model</p>206        207        <h2>Available Endpoints:</h2>208        <div class="endpoint"><strong>GET /health</strong> - Health check</div>209        <div class="endpoint"><strong>POST /review</strong> - Submit code diff for review</div>210        <div class="endpoint"><strong>GET /docs</strong> - Interactive API documentation</div>211        212        <h2>Quick Test:</h2>213        <p><a href="/health">Test Health Endpoint</a></p>214        <p><a href="/docs">View API Documentation</a></p>215        216        <h2>Status:</h2>217        <ul>218            <li>Mode: Hugging Face API</li>219            <li>AI Model: GPT-2</li>220            <li>Response Time: ~2-5 seconds</li>221        </ul>222    </body>223    </html>224    """225 226@app.get("/health")227async def health_check():228    return {"status": "ok", "api_configured": HF_API_KEY is not None}