main_lightweight.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import logging
import json
import re

# Ultra-lightweight version with minimal AI
app = FastAPI(
    title="AI Code Review Service",
    description="An API to get AI-powered code reviews for pull request diffs.",
    version="1.0.0",
)

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Try to load a very small model, fall back to mock if it fails
model = None
tokenizer = None

def load_simple_model():
    """Try to load the smallest possible model."""
    global model, tokenizer
    try:
        from transformers import AutoTokenizer, AutoModelForCausalLM
        
        # Use the smallest possible model
        model_name = "distilgpt2"  # Much smaller than TinyLlama
        logger.info("Loading lightweight model: %s", model_name)
        
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
        )
        logger.info("Model loaded successfully")
        return True
    except Exception as e:
        logger.warning("Failed to load AI model: %s. Using mock responses.", str(e))
        return False

# Try to load model on startup
model_loaded = load_simple_model()

class DiffRequest(BaseModel):
    diff: str

class ReviewComment(BaseModel):
    file_path: str
    line_number: int
    comment_text: str

class ReviewResponse(BaseModel):
    comments: list[ReviewComment]

@app.get("/health")
def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "service": "AI Code Review Service",
        "model_loaded": model_loaded,
        "model_name": "distilgpt2" if model_loaded else "mock",
        "device": "cuda" if torch.cuda.is_available() else "cpu"
    }

def simple_ai_review(diff: str):
    """Very simple AI review using the lightweight model."""
    if not model_loaded or not model or not tokenizer:
        return None
        
    try:
        # Very simple prompt
        prompt = f"Review this code change and suggest improvements:\n{diff[:200]}\nSuggestion:"
        
        inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=256, truncation=True)
        
        # Very conservative generation
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=50,  # Very short response
                do_sample=False,
                num_return_sequences=1,
                pad_token_id=tokenizer.eos_token_id,
                use_cache=True
            )
        
        response = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
        return response.strip()
    except Exception as e:
        logger.warning("AI generation failed: %s", str(e))
        return None

@app.post("/review", response_model=ReviewResponse)
def review_diff(request: DiffRequest):
    """Review endpoint with fallback to mock data."""
    logger.info("Received diff for review (length: %d chars)", len(request.diff))
    
    # Try AI first, fall back to mock
    ai_suggestion = None
    if model_loaded:
        ai_suggestion = simple_ai_review(request.diff)
    
    if ai_suggestion:
        # Use AI suggestion
        comments = [{
            "file_path": "reviewed_file.py",
            "line_number": 1,
            "comment_text": ai_suggestion
        }]
        logger.info("Returning AI-generated review")
    else:
        # Fall back to mock comments
        comments = [
            {
                "file_path": "example.py",
                "line_number": 1,
                "comment_text": "Consider adding error handling and input validation."
            },
            {
                "file_path": "example.py",
                "line_number": 5,
                "comment_text": "This function could benefit from better documentation."
            }
        ]
        logger.info("Returning mock review comments")
    
    return ReviewResponse(comments=comments)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)