Update server.py

server.py

@@ -1,5 +1,5 @@
 from typing import Optional, Tuple
-from fastapi import FastAPI, UploadFile, File
+from fastapi import FastAPI, UploadFile, File, Form
 from fastapi.responses import FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from PIL import Image, ExifTags
@@ -9,10 +9,18 @@ import httpx
 import os
 import base64
 import json
+import asyncio
+import cv2
+import tempfile
+import fitz  # pymupdf
+import pypdf
 
 # ---------------- API KEYS ----------------
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+HF_API_KEY = os.getenv("HF_API_KEY")
 GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
+ROBERTA_FAKE_NEWS_URL = "https://router.huggingface.co/hf-inference/models/hamzab/roberta-fake-news-classification"
+ROBERTA_AI_TEXT_URL = "https://router.huggingface.co/hf-inference/models/openai-community/roberta-base-openai-detector"
 
 # ---------------- APP ----------------
 app = FastAPI()
@@ -53,15 +61,39 @@ def calculate_metadata_risk(image: Image.Image):
 
 
 def fusion_score(model_score: float, metadata_risk: float):
-    final = 0.9 * model_score + 0.2 * metadata_risk
+    final = 0.9 * model_score + 0.1 * metadata_risk
     authenticity = (1 - final) * 100
     fake = final * 100
     return authenticity, fake
 
 
+def normalize_output(label_prob_dict: dict) -> float:
+    FAKE_KEYWORDS = ["fake", "ai", "generated", "manipulated", "deepfake", "artificial", "synthetic", "machine"]
+    REAL_KEYWORDS = ["real", "authentic", "genuine", "human", "original"]
+
+    fake_score = 0.0
+    uncertain_score = 0.0
+
+    for label, prob in label_prob_dict.items():
+        label_lower = label.lower()
+        if any(k in label_lower for k in FAKE_KEYWORDS):
+            fake_score += prob
+        elif any(k in label_lower for k in REAL_KEYWORDS):
+            pass
+        else:
+            uncertain_score += prob
+
+    fake_score += 0.4 * uncertain_score
+    return min(fake_score, 1.0)
+
+
+def make_confidence(authenticity, fake):
+    diff = abs(authenticity - fake)
+    return "low" if diff < 20 else "medium" if diff < 40 else "high"
+
+
 # ==============================
-# GROQ VISION
-# now returns (score, reasoning) tuple
+# GROQ VISION (images)
 # ==============================
 
 async def call_groq_vision(contents: bytes) -> Tuple[Optional[float], str]:
@@ -121,7 +153,6 @@ fake_probability must be between 0.0 (definitely real) and 1.0 (definitely AI/fa
             data = response.json()
             text = data["choices"][0]["message"]["content"]
             print(f"Groq vision response: {text}")
-
             clean = text.strip().replace("```json", "").replace("```", "")
             result = json.loads(clean)
             return float(result["fake_probability"]), result.get("reasoning", "")
@@ -131,6 +162,91 @@ fake_probability must be between 0.0 (definitely real) and 1.0 (definitely AI/fa
         return None, ""
 
 
+# ==============================
+# GROQ TEXT (for AI writing detection)
+# ==============================
+
+async def call_groq_text(text: str) -> Tuple[Optional[float], str]:
+    if not GROQ_API_KEY:
+        return None, ""
+    try:
+        payload = {
+            "model": "llama-3.3-70b-versatile",
+            "messages": [
+                {
+                    "role": "user",
+                    "content": f"""You are a forensic text analyst. Analyze the following text and determine if it is AI-generated or written by a human. Also check if it could be a forged government document or fake news.
+
+Look for:
+- Overly formal or repetitive sentence structure typical of LLMs
+- Lack of personal voice or human inconsistencies
+- Suspiciously perfect grammar with no natural errors
+- Generic phrasing commonly used by AI models
+- For government documents: inconsistent terminology, wrong formats, suspicious clauses
+- For news: sensational language, lack of credible sources, misleading framing
+
+Text to analyze:
+\"\"\"
+{text[:4000]}
+\"\"\"
+
+Respond ONLY in this exact JSON format, nothing else:
+{{"fake_probability": 0.0, "reasoning": "brief reason"}}
+
+fake_probability must be between 0.0 (definitely human/authentic) and 1.0 (definitely AI-generated/forged)."""
+                }
+            ],
+            "max_tokens": 200,
+            "temperature": 0.1
+        }
+
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                GROQ_API_URL,
+                headers={
+                    "Authorization": f"Bearer {GROQ_API_KEY}",
+                    "Content-Type": "application/json"
+                },
+                json=payload
+            )
+            response.raise_for_status()
+            data = response.json()
+            text_response = data["choices"][0]["message"]["content"]
+            print(f"Groq text response: {text_response}")
+            clean = text_response.strip().replace("```json", "").replace("```", "")
+            result = json.loads(clean)
+            return float(result["fake_probability"]), result.get("reasoning", "")
+
+    except Exception as e:
+        print(f"Groq text failed: {e}")
+        return None, ""
+
+
+# ==============================
+# ROBERTA (for fake news + AI text)
+# ==============================
+
+async def call_roberta(url: str, text: str, name: str) -> Optional[float]:
+    if not HF_API_KEY:
+        print(f"No HF_API_KEY, skipping {name}")
+        return None
+    try:
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                url,
+                headers={"Authorization": f"Bearer {HF_API_KEY}"},
+                json={"inputs": text[:512]}
+            )
+            response.raise_for_status()
+            data = response.json()
+            print(f"{name} response: {data}")
+            label_prob_dict = {item["label"]: item["score"] for item in data[0]}
+            return normalize_output(label_prob_dict)
+    except Exception as e:
+        print(f"{name} failed: {e}")
+        return None
+
+
 # ==============================
 # ANALYZERS
 # ==============================
@@ -138,7 +254,11 @@ fake_probability must be between 0.0 (definitely real) and 1.0 (definitely AI/fa
 async def analyze_image(contents: bytes, content_type: str = "image/jpeg"):
     image = Image.open(io.BytesIO(contents)).convert("RGB")
 
-    score, reasoning = await call_groq_vision(contents)
+    if len(contents) > 20 * 1024 * 1024:
+        print("Image too large for Groq")
+        score, reasoning = None, "Image too large for analysis"
+    else:
+        score, reasoning = await call_groq_vision(contents)
 
     combined_model_score = score if score is not None else 0.5
     models_used = ["Groq_Llama4"] if score is not None else []
@@ -146,14 +266,11 @@ async def analyze_image(contents: bytes, content_type: str = "image/jpeg"):
     metadata_risk = calculate_metadata_risk(image)
     authenticity, fake = fusion_score(combined_model_score, metadata_risk)
 
-    diff = abs(authenticity - fake)
-    confidence_level = "low" if diff < 20 else "medium" if diff < 40 else "high"
-
     return {
         "type": "image",
         "authenticity": round(authenticity, 2),
         "fake": round(fake, 2),
-        "confidence_level": confidence_level,
+        "confidence_level": make_confidence(authenticity, fake),
         "models_used": models_used,
         "details": {
             "groq_score": round(score, 4) if score is not None else "unavailable",
@@ -163,13 +280,7 @@ async def analyze_image(contents: bytes, content_type: str = "image/jpeg"):
     }
 
 
-import cv2
-import numpy as np
-import tempfile
-import asyncio
-
 async def analyze_video(contents: bytes):
-    # write to temp file since cv2 needs a path
     with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as f:
         f.write(contents)
         tmp_path = f.name
@@ -180,7 +291,6 @@ async def analyze_video(contents: bytes):
         fps = cap.get(cv2.CAP_PROP_FPS)
         duration = round(frame_count / fps, 1) if fps > 0 else 0
 
-        # sample 5 frames evenly across the video
         sample_indices = [int(frame_count * i / 5) for i in range(5)]
         frames = []
 
@@ -188,7 +298,6 @@ async def analyze_video(contents: bytes):
             cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
             ret, frame = cap.read()
             if ret:
-                # convert frame to jpeg bytes
                 pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                 buf = io.BytesIO()
                 pil_img.save(buf, format="JPEG", quality=85)
@@ -213,7 +322,6 @@ async def analyze_video(contents: bytes):
                 }
             }
 
-        # analyze each frame with Groq — 2s delay between calls to avoid rate limit
         scores = []
         reasonings = []
         for i, frame_bytes in enumerate(frames):
@@ -223,28 +331,20 @@ async def analyze_video(contents: bytes):
                 scores.append(score)
                 reasonings.append(f"Frame {i+1}: {reasoning}")
             if i < len(frames) - 1:
-                await asyncio.sleep(2)  # avoid rate limit
+                await asyncio.sleep(2)
 
-        if scores:
-            combined_model_score = sum(scores) / len(scores)
-            models_used = ["Groq_Llama4"]
-            groq_reasoning = " | ".join(reasonings)
-        else:
-            combined_model_score = 0.5
-            models_used = []
-            groq_reasoning = "All frame analyses failed."
+        combined_model_score = sum(scores) / len(scores) if scores else 0.5
+        models_used = ["Groq_Llama4"] if scores else []
+        groq_reasoning = " | ".join(reasonings) if reasonings else "All frame analyses failed."
 
         authenticity = round((1 - combined_model_score) * 100, 2)
         fake = round(combined_model_score * 100, 2)
 
-        diff = abs(authenticity - fake)
-        confidence_level = "low" if diff < 20 else "medium" if diff < 40 else "high"
-
         return {
             "type": "video",
             "authenticity": authenticity,
             "fake": fake,
-            "confidence_level": confidence_level,
+            "confidence_level": make_confidence(authenticity, fake),
             "models_used": models_used,
             "details": {
                 "groq_score": round(combined_model_score, 4),
@@ -275,21 +375,150 @@ async def analyze_video(contents: bytes):
         }
 
 
+async def analyze_text(text: str):
+    # all 3 in parallel
+    results = await asyncio.gather(
+        call_roberta(ROBERTA_FAKE_NEWS_URL, text, "RoBERTa_FakeNews"),
+        call_roberta(ROBERTA_AI_TEXT_URL, text, "RoBERTa_AIDetector"),
+        call_groq_text(text)
+    )
+
+    score1 = results[0]
+    score2 = results[1]
+    score3, reasoning = results[2]
+
+    scores = [(s, n) for s, n in [
+        (score1, "RoBERTa_FakeNews"),
+        (score2, "RoBERTa_AIDetector"),
+        (score3, "Groq_Llama3")
+    ] if s is not None]
+
+    combined = sum(s for s, _ in scores) / len(scores) if scores else 0.5
+    models_used = [n for _, n in scores]
+
+    authenticity = round((1 - combined) * 100, 2)
+    fake = round(combined * 100, 2)
+
+    return {
+        "type": "text",
+        "authenticity": authenticity,
+        "fake": fake,
+        "confidence_level": make_confidence(authenticity, fake),
+        "models_used": models_used,
+        "details": {
+            "groq_score": round(score3, 4) if score3 is not None else "unavailable",
+            "roberta_fakenews_score": round(score1, 4) if score1 is not None else "unavailable",
+            "roberta_aidetector_score": round(score2, 4) if score2 is not None else "unavailable",
+            "groq_reasoning": reasoning,
+            "metadata_risk": 0.0,
+        }
+    }
+
+
+async def analyze_pdf(contents: bytes):
+    scores = []
+    reasonings = []
+
+    try:
+        # extract text
+        reader = pypdf.PdfReader(io.BytesIO(contents))
+        full_text = ""
+        for page in reader.pages:
+            full_text += page.extract_text() or ""
+
+        if full_text.strip():
+            print(f"Extracted {len(full_text)} chars from PDF")
+            text_results = await asyncio.gather(
+                call_roberta(ROBERTA_FAKE_NEWS_URL, full_text, "RoBERTa_FakeNews"),
+                call_roberta(ROBERTA_AI_TEXT_URL, full_text, "RoBERTa_AIDetector"),
+                call_groq_text(full_text)
+            )
+            s1 = text_results[0]
+            s2 = text_results[1]
+            s3, text_reasoning = text_results[2]
+
+            if s1 is not None:
+                scores.append(s1)
+                reasonings.append(f"RoBERTa FakeNews: {round(s1*100)}% fake")
+            if s2 is not None:
+                scores.append(s2)
+                reasonings.append(f"RoBERTa AI Detector: {round(s2*100)}% AI-generated")
+            if s3 is not None:
+                scores.append(s3)
+                reasonings.append(f"Groq text: {text_reasoning}")
+
+        # extract and analyze images inside PDF
+        doc = fitz.open(stream=contents, filetype="pdf")
+        image_count = 0
+        for page in doc:
+            for img in page.get_images():
+                if image_count >= 3:
+                    break
+                xref = img[0]
+                base_image = doc.extract_image(xref)
+                img_bytes = base_image["image"]
+                await asyncio.sleep(2)
+                img_score, img_reasoning = await call_groq_vision(img_bytes)
+                if img_score is not None:
+                    scores.append(img_score)
+                    reasonings.append(f"Image {image_count+1}: {img_reasoning}")
+                image_count += 1
+        doc.close()
+
+    except Exception as e:
+        print(f"PDF analysis error: {e}")
+
+    combined = sum(scores) / len(scores) if scores else 0.5
+    models_used = ["RoBERTa_FakeNews", "RoBERTa_AIDetector", "Groq_Llama3+Vision"] if scores else []
+
+    authenticity = round((1 - combined) * 100, 2)
+    fake = round(combined * 100, 2)
+
+    return {
+        "type": "pdf",
+        "authenticity": authenticity,
+        "fake": fake,
+        "confidence_level": make_confidence(authenticity, fake),
+        "models_used": models_used,
+        "details": {
+            "groq_score": "see breakdown",
+            "groq_reasoning": " | ".join(reasonings) if reasonings else "No content extracted",
+            "metadata_risk": 0.0,
+        }
+    }
+
+
 # ==============================
 # ROUTER
 # ==============================
 
 @app.post("/analyze")
-async def analyze(file: UploadFile = File(...)):
+async def analyze(
+    file: Optional[UploadFile] = File(None),
+    text: Optional[str] = Form(None)
+):
+    # plain text input
+    if text and not file:
+        result = await analyze_text(text)
+        result["sha256"] = hashlib.sha256(text.encode()).hexdigest()
+        return result
+
+    if not file:
+        return {"error": "No file or text provided"}
+
     contents = await file.read()
     sha256 = calculate_sha256(contents)
 
     if file.content_type.startswith("image/"):
         result = await analyze_image(contents, file.content_type)
     elif file.content_type.startswith("video/"):
-        result = await analyze_video(contents)  # ← await added
+        result = await analyze_video(contents)
+    elif file.content_type == "application/pdf":
+        result = await analyze_pdf(contents)
     else:
         return {"error": "Unsupported file type"}
 
     result["sha256"] = sha256
-    return result
+    return result
+
+