gemini_classifier.py

import os
import json
import google.generativeai as genai
from dotenv import load_dotenv

load_dotenv()

class GeminiClassifier:
    def __init__(self):
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            raise ValueError("GEMINI_API_KEY not found in .env file")
        
        genai.configure(api_key=api_key)
        
        # We use strict JSON generation for the classification
        self.generation_config = {
            "temperature": 0.1,  # Low temperature for highly deterministic classification
            "response_mime_type": "application/json",
        }
        
        # Using Gemini 1.5 Pro or Flash
        self.model = genai.GenerativeModel(
            model_name="gemini-1.5-flash",
            generation_config=self.generation_config,
        )

    def analyze_text(self, text):
        """
        Zero-shot prompt to classify Hinglish scam texts.
        Returns a structured JSON with fraud score, category, and red flags.
        """
        if not text or len(text) < 5:
            return {
                "fraud_score": 0,
                "scam_category": "Unclear / Too Short",
                "red_flags": [],
                "summary": "The provided text is too short or empty. Provide a clear screenshot."
            }

        prompt = f"""
You are an expert cybersecurity analyst at 'RealityGuard India', specialized in analyzing WhatsApp and SMS messages for scams targeting Indian citizens.
Your task is to analyze the following Hinglish (Hindi + English) text extracted from a message.

SCAM CRITERIA TO LOOK FOR (Common in India):
1. **Urgency/Fear:** Threats of electricity disconnection, bank account block, SBI/HDFC KYC updates, police arrest.
2. **Greed/Lottery:** KBC lottery win, Jio tower installation, work-from-home YouTube likes jobs, free cryptocurrency.
3. **Suspicious Links:** Bit.ly links, .apk downloads, unofficial portals mimicking Indian banks or gov.
4. **Payment Requests:** Unknown UPI VPAs (e.g., @ybl, @okicici), requesting ₹10 or ₹5000 as "processing fee".

Analyze the text and return ONLY a valid JSON object with the following schema:
{{
  "fraud_score": <integer from 0 to 100, where 100 is definitely a scam>,
  "scam_category": "<string, e.g., 'Phishing', 'Lottery Scam', 'Electricity Bill Scam', 'Job Fraud', or 'Safe'>",
  "red_flags": [
    "<string: explicitly list the problematic parts of the text, e.g., 'Mentions urgency for KYC', 'Contains unofficial APK link'>"
  ],
  "summary": "<string: A 1-2 sentence Hinglish or English explanation for the user on what this message is and what they should do.>"
}}

TEXT TO ANALYZE:
"{text}"
        """

        try:
            response = self.model.generate_content(prompt)
            # Parse the JSON response explicitly
            result = json.loads(response.text)
            return result
        except Exception as e:
            # Fallback error structure
            return {
                "fraud_score": 0,
                "scam_category": "Analysis Failed",
                "red_flags": [],
                "summary": f"Failed to analyze text using Gemini API. Error: {str(e)}"
            }

if __name__ == "__main__":
    # Test
    classifier = GeminiClassifier()
    print("Gemini Classifier initialized.")