gemini_ai.py

import json
import os
from google import genai
from google.genai import types
from pydantic import BaseModel
from typing import Optional

# IMPORTANT: KEEP THIS COMMENT
# Follow these instructions when using this blueprint:
# - Note that the newest Gemini model series is "gemini-2.5-flash" or gemini-2.5-pro"
# - The SDK was recently renamed from google-generativeai to google-genai

client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))


class IssueDiagnosis(BaseModel):
    summary: str
    affected_files: list[str]
    root_cause: str
    complexity: str


class CodePlan(BaseModel):
    approach: str
    files_to_modify: list[str]
    estimated_changes: int
    risks: list[str]


class CodePatch(BaseModel):
    file_path: str
    original_code: str
    modified_code: str
    explanation: str


def analyze_issue(issue_text: str, repo_context: str) -> IssueDiagnosis:
    """Analyze a GitHub issue and diagnose the problem."""
    system_prompt = """You are an expert code analyst. Analyze the GitHub issue and repository context.
    Identify the affected files, root cause, and complexity level (low/medium/high).
    Return structured JSON with: summary, affected_files (list), root_cause, complexity."""
    
    prompt = f"""Repository Context:
{repo_context}

GitHub Issue:
{issue_text}

Analyze this issue and provide a diagnosis."""

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[types.Content(role="user", parts=[types.Part(text=prompt)])],
            config=types.GenerateContentConfig(
                system_instruction=system_prompt,
                response_mime_type="application/json",
                response_schema=IssueDiagnosis,
            ),
        )
        
        if hasattr(response, 'parsed') and response.parsed:
            return response.parsed
        elif response.text:
            data = json.loads(response.text)
            return IssueDiagnosis(**data)
        else:
            raise ValueError("Empty response from model")
    except Exception as e:
        raise Exception(f"Failed to analyze issue: {e}")


def generate_fix_plan(diagnosis: IssueDiagnosis, repo_structure: str) -> CodePlan:
    """Generate a plan to fix the issue."""
    system_prompt = """You are an expert software architect. Create a detailed plan to fix the issue.
    Specify approach, files to modify, estimated number of changes, and potential risks.
    Return structured JSON."""
    
    prompt = f"""Diagnosis:
{diagnosis.model_dump_json(indent=2)}

Repository Structure:
{repo_structure}

Create a detailed fix plan."""

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[types.Content(role="user", parts=[types.Part(text=prompt)])],
            config=types.GenerateContentConfig(
                system_instruction=system_prompt,
                response_mime_type="application/json",
                response_schema=CodePlan,
            ),
        )
        
        if hasattr(response, 'parsed') and response.parsed:
            return response.parsed
        elif response.text:
            data = json.loads(response.text)
            return CodePlan(**data)
        else:
            raise ValueError("Empty response from model")
    except Exception as e:
        raise Exception(f"Failed to generate plan: {e}")


def generate_code_patch(plan: CodePlan, file_content: str, file_path: str, diagnosis: IssueDiagnosis) -> CodePatch:
    """Generate actual code patch for a specific file."""
    system_prompt = """You are an expert programmer. Generate a precise code patch to fix the issue.
    Provide the original code section, modified code, and explanation.
    Follow best practices and maintain code style.
    Return structured JSON."""
    
    prompt = f"""Fix Plan:
{plan.approach}

File Path: {file_path}

Current File Content:
{file_content}

Issue Root Cause: {diagnosis.root_cause}

Generate a code patch to fix this issue. Only modify what's necessary."""

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[types.Content(role="user", parts=[types.Part(text=prompt)])],
            config=types.GenerateContentConfig(
                system_instruction=system_prompt,
                response_mime_type="application/json",
                response_schema=CodePatch,
            ),
        )
        
        if hasattr(response, 'parsed') and response.parsed:
            return response.parsed
        elif response.text:
            data = json.loads(response.text)
            return CodePatch(**data)
        else:
            raise ValueError("Empty response from model")
    except Exception as e:
        raise Exception(f"Failed to generate patch: {e}")


def analyze_pr_quality(pr_title: str, pr_description: str, diff_content: str, file_changes: int) -> dict:
    """Analyze PR quality and provide a quality score."""
    prompt = f"""Analyze this Pull Request for quality:

Title: {pr_title}
Description: {pr_description}
Files Changed: {file_changes}

Diff Content (first 2000 chars):
{diff_content[:2000]}

Evaluate on these criteria (score 0-10 each):
1. Code quality and best practices
2. Clear problem statement and solution
3. Appropriate scope (not too large)
4. Good commit message and PR description

Return JSON with: {{
    "code_quality_score": <0-10>,
    "description_score": <0-10>,
    "scope_score": <0-10>,
    "documentation_score": <0-10>,
    "overall_score": <0-10>,
    "strengths": [<list of strengths>],
    "weaknesses": [<list of weaknesses>],
    "recommendation": "<merge/needs_work/reject>"
}}"""

    try:
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=prompt
        )
        
        if response.text:
            # Parse JSON from response
            text = response.text.strip()
            if text.startswith("```json"):
                text = text[7:]
            if text.endswith("```"):
                text = text[:-3]
            return json.loads(text.strip())
        else:
            return {
                "code_quality_score": 5,
                "description_score": 5,
                "scope_score": 5,
                "documentation_score": 5,
                "overall_score": 5,
                "strengths": [],
                "weaknesses": ["Unable to analyze"],
                "recommendation": "needs_work"
            }
    except Exception as e:
        print(f"Error analyzing PR quality: {e}")
        return {
            "code_quality_score": 5,
            "description_score": 5,
            "scope_score": 5,
            "documentation_score": 5,
            "overall_score": 5,
            "strengths": [],
            "weaknesses": [f"Analysis error: {str(e)}"],
            "recommendation": "needs_work"
        }