Initial commit

Dockerfile

@@ -0,0 +1,20 @@
+
+# Use the official Python 3.9 slim image
+FROM python:3.9-slim
+
+# Set the working directory inside the container
+WORKDIR /code
+
+# Copy the requirements file into the container
+COPY ./requirements.txt /code/requirements.txt
+
+# Install the Python dependencies
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Copy the main application file into the container
+COPY ./main.py /code/main.py
+
+# Command to run the FastAPI server with Uvicorn
+# We use --host 0.0.0.0 to make it accessible from outside the container
+# and --port 7860 as this is the standard port Hugging Face Spaces expects
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,195 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import re
+import json
+
+# ----------------------------
+# 1. Configuration
+# ----------------------------
+
+# Define the model we want to use. 
+# We use a 4-bit quantized version ("4bit") for efficiency.
+MODEL_NAME = "deepseek-ai/deepseek-coder-6.7b-instruct"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ----------------------------
+# 2. FastAPI App Initialization
+# ----------------------------
+
+app = FastAPI(
+    title="AI Code Review Service",
+    description="An API to get AI-powered code reviews for pull request diffs.",
+    version="1.0.0",
+)
+
+# ----------------------------
+# 3. AI Model Loading
+# ----------------------------
+
+# Use a global variable to hold the model and tokenizer
+# This is lazy-loaded on the first request to speed up server startup.
+model = None
+tokenizer = None
+
+def load_model():
+    """Loads the model and tokenizer into memory."""
+    global model, tokenizer
+    if model is None:
+        print(f"Loading model: {MODEL_NAME} on device: {DEVICE}...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+        
+        # Load the model with 4-bit quantization to save memory
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            trust_remote_code=True,
+            torch_dtype=torch.bfloat16,
+            load_in_4bit=True,
+        )
+        print("Model loaded successfully.")
+
+@app.on_event("startup")
+async def startup_event():
+    """
+    On server startup, we trigger the model loading.
+    This makes the first API call after startup faster.
+    """
+    print("Server starting up...")
+    load_model()
+
+# ----------------------------
+# 4. API Request/Response Models
+# ----------------------------
+
+class ReviewRequest(BaseModel):
+    """The request body for the /review endpoint."""
+    diff: str
+
+class ReviewComment(BaseModel):
+    """A single review comment."""
+    file_path: str
+    line_number: int
+    comment_text: str
+
+class ReviewResponse(BaseModel):
+    """The response body for the /review endpoint."""
+    comments: list[ReviewComment]
+
+# ----------------------------
+# 5. The AI Review Logic
+# ----------------------------
+
+def run_ai_inference(diff: str) -> str:
+    """
+    Runs the AI model to get the review.
+    """
+    if not model or not tokenizer:
+        raise RuntimeError("Model is not loaded.")
+
+    # This is the prompt engineering part. We create a clear instruction
+    # for the model, telling it exactly what to do and what format to output.
+    messages = [
+        {
+            "role": "system",
+            "content": """
+You are an expert code reviewer. Your task is to analyze a pull request diff and provide constructive feedback.
+Analyze the provided diff and identify potential issues, suggest improvements, or point out good practices.
+Your feedback should be in the form of review comments.
+
+IMPORTANT: Respond with a JSON array of comment objects. Each object must have three fields: 'file_path', 'line_number', and 'comment_text'.
+The 'file_path' should be the full path of the file being changed.
+The 'line_number' must be an integer corresponding to the line number in the *new* version of the file where the comment applies.
+The 'comment_text' should be your concise and clear review comment.
+
+Example response format:
+[
+    {
+        "file_path": "src/utils/helpers.py",
+        "line_number": 42,
+        "comment_text": "This function could be simplified by using a list comprehension."
+    },
+    {
+        "file_path": "README.md",
+        "line_number": 12,
+        "comment_text": "There is a typo in this sentence."
+    }
+]
+
+Do not add any introductory text or explanations outside of the JSON array.
+"""
+        },
+        {
+            "role": "user",
+            "content": f"Here is the diff to review:\n\n```diff\n{diff}\n```"
+        }
+    ]
+
+    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
+    
+    # Generate the response from the model
+    outputs = model.generate(inputs, max_new_tokens=1024, do_sample=False, top_k=50, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
+    
+    # Decode the output and clean it up
+    response_text = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+    return response_text.strip()
+
+def parse_ai_response(response_text: str) -> list[ReviewComment]:
+    """
+    Parses the raw text from the AI to extract the JSON array.
+    This function is robust against the AI adding extra text before or after the JSON.
+    """
+    print(f"Raw AI Response:\n---\n{response_text}\n---")
+    
+    # Find the start and end of the JSON array
+    json_match = re.search(r'\[.*\]', response_text, re.DOTALL)
+    if not json_match:
+        print("Warning: Could not find a JSON array in the AI response.")
+        return []
+
+    json_string = json_match.group(0)
+    
+    try:
+        comments_data = json.loads(json_string)
+        # Validate the structure of the parsed data
+        validated_comments = [ReviewComment(**item) for item in comments_data]
+        return validated_comments
+    except (json.JSONDecodeError, TypeError, KeyError) as e:
+        print(f"Error parsing JSON from AI response: {e}")
+        print(f"Invalid JSON string: {json_string}")
+        return []
+
+# ----------------------------
+# 6. The API Endpoint
+# ----------------------------
+
+@app.post("/review", response_model=ReviewResponse)
+async def get_code_review(request: ReviewRequest):
+    """
+    Receives a code diff, gets a review from the AI model,
+    and returns structured review comments.
+    """
+    if not request.diff:
+        raise HTTPException(status_code=400, detail="Diff content cannot be empty.")
+
+    try:
+        # 1. Run the AI model
+        ai_response_text = run_ai_inference(request.diff)
+        
+        # 2. Parse the AI's response into structured objects
+        parsed_comments = parse_ai_response(ai_response_text)
+        
+        return ReviewResponse(comments=parsed_comments)
+
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+        raise HTTPException(status_code=500, detail="An internal error occurred while processing the review.")
+
+# ----------------------------
+# 7. Health Check Endpoint
+# ----------------------------
+
+@app.get("/health")
+async def health_check():
+    """A simple endpoint to confirm the server is running."""
+    return {"status": "ok", "model_loaded": model is not None}

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+pydantic
+torch
+transformers
+accelerate
+bitsandbytes