Update app/services/reviewer_service.py

app/services/reviewer_service.py

@@ -1,58 +1,116 @@
 import json
 import logging
 import re
+import hashlib
+import sqlite3
+import threading
 from typing import List, Dict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from app.core.model_loader import llm_engine
 
 logger = logging.getLogger(__name__)
 
+DB_PATH = "review_cache.db"
 
 class AIReviewerService:
     def __init__(self):
-        # Keep worker count under the Gemini free-tier RPM limit (15/min).
-        # 3 workers gives real parallelism without immediately tripping 429s,
-        # since llm_engine already has its own RPM guard + backoff on top.
-        self.max_concurrent_batches = 50
+        self.max_concurrent_batches = 3
         self.batch_size = 10
+        self.db_lock = threading.Lock()
+        self._init_db()
+
+    def _init_db(self):
+        """Initializes SQLite cache table."""
+        with self.db_lock:
+            with sqlite3.connect(DB_PATH) as conn:
+                cursor = conn.cursor()
+                cursor.execute("""
+                    CREATE TABLE IF NOT EXISTS review_cache (
+                        file_hash TEXT PRIMARY KEY,
+                        result_json TEXT
+                    )
+                """)
+                conn.commit()
+
+    def _generate_file_hash(self, file_obj) -> str:
+        """Generates a unique SHA256 digest based on path and content."""
+        content = file_obj.content or ""
+        raw_payload = f"{file_obj.fileName}:{content}"
+        return hashlib.sha256(raw_payload.encode("utf-8")).hexdigest()
+
+    def _get_cached_result(self, file_hash: str) -> dict:
+        """Thread-safe SQLite fetch."""
+        with self.db_lock:
+            with sqlite3.connect(DB_PATH) as conn:
+                cursor = conn.cursor()
+                cursor.execute("SELECT result_json FROM review_cache WHERE file_hash = ?", (file_hash,))
+                row = cursor.fetchone()
+                if row:
+                    try:
+                        return json.loads(row[0])
+                    except json.JSONDecodeError:
+                        return None
+        return None
+
+    def _save_cached_results(self, hash_result_pairs: List[tuple]):
+        """Thread-safe batch insert into SQLite."""
+        if not hash_result_pairs:
+            return
+        with self.db_lock:
+            with sqlite3.connect(DB_PATH) as conn:
+                cursor = conn.cursor()
+                cursor.executemany(
+                    "INSERT OR REPLACE INTO review_cache (file_hash, result_json) VALUES (?, ?)",
+                    [(h, json.dumps(res)) for h, res in hash_result_pairs]
+                )
+                conn.commit()
 
     def review_batch_code(self, files: list) -> list:
-        """
-        Splits files into batches and reviews them CONCURRENTLY using a thread
-        pool, instead of the old sequential for-loop. Each batch still makes
-        one blocking call to llm_engine.generate(), but multiple batches are
-        now in flight at once.
-        """
-        batches = [
-            files[i : i + self.batch_size]
-            for i in range(0, len(files), self.batch_size)
-        ]
-
-        if not batches:
+        if not files:
             return []
 
-        results: list = []
+        results = [None] * len(files)
+        uncached_files = []
+        uncached_indices = []
+
+        # 1. Check SQLite cache
+        for idx, f in enumerate(files):
+            file_hash = self._generate_file_hash(f)
+            cached_res = self._get_cached_result(file_hash)
+            if cached_res:
+                results[idx] = cached_res
+            else:
+                uncached_files.append(f)
+                uncached_indices.append(idx)
+
+        if not uncached_files:
+            logger.info(f"⚡ Served all {len(files)} files instantly from SQLite cache.")
+            return results
+
+        logger.info(f"SQLite cache hit: {len(files) - len(uncached_files)}/{len(files)} files. Reviewing {len(uncached_files)} uncached files.")
+
+        # 2. Batch process uncached files via Gemini
+        batches = [
+            uncached_files[i : i + self.batch_size]
+            for i in range(0, len(uncached_files), self.batch_size)
+        ]
 
+        new_results = []
         with ThreadPoolExecutor(max_workers=self.max_concurrent_batches) as pool:
             future_to_batch = {
                 pool.submit(self._process_batch, batch): batch for batch in batches
             }
-
             for future in as_completed(future_to_batch):
                 batch = future_to_batch[future]
                 try:
                     batch_results = future.result()
-                    results.extend(batch_results)
+                    new_results.extend(batch_results)
                 except Exception as e:
-                    # If a whole batch blew up unexpectedly (shouldn't normally
-                    # happen since _process_batch catches its own errors), fall
-                    # back to per-file error entries so the response shape
-                    # stays consistent for the Android client.
-                    logger.error(f"Unhandled batch failure: {e}")
-                    for f in batch:
-                        results.append(
+                    logger.error(f"Batch failure: {e}")
+                    for f_item in batch:
+                        new_results.append(
                             {
-                                "fileName": f.fileName,
+                                "fileName": f_item.fileName,
                                 "vulnerabilities": [],
                                 "improvement_suggestions": [],
                                 "metrics": {"complexity": 0, "maintainability": 0},
@@ -60,36 +118,49 @@ class AIReviewerService:
                             }
                         )
 
+        new_results_map = {item.get("fileName"): item for item in new_results}
+        to_cache = []
+
+        # 3. Assemble final response list and write to SQLite
+        for idx, f in zip(uncached_indices, uncached_files):
+            res = new_results_map.get(
+                f.fileName,
+                {
+                    "fileName": f.fileName,
+                    "vulnerabilities": [],
+                    "improvement_suggestions": [],
+                    "metrics": {"complexity": 0, "maintainability": 0},
+                    "error": "analysis_failed",
+                },
+            )
+
+            if "error" not in res:
+                file_hash = self._generate_file_hash(f)
+                to_cache.append((file_hash, res))
+
+            results[idx] = res
+
+        self._save_cached_results(to_cache)
         return results
 
     def _process_batch(self, batch: list) -> list:
-        """
-        Processes a single batch of files: builds the prompt, calls Gemini,
-        parses the response, and normalizes each file's result. This is the
-        unit of work that now runs concurrently across threads.
-        """
         combined_code = ""
         file_names = []
-
         for f in batch:
             raw_content = f.content or ""
-            # Limit to 6k chars per file to stay within context window safely
-            minified_content = self._minify_code(raw_content[:6000], f.fileName)
-
+            minified_content = self._minify_code(raw_content[:5000], f.fileName)
             combined_code += f"\n--- FILE: {f.fileName} ---\n{minified_content}\n"
             file_names.append(f.fileName)
 
-        prompt = f"""
-Analyze the following {len(batch)} source code files for security vulnerabilities and code quality.
+        prompt = f"""Analyze these {len(batch)} source files for vulnerabilities and quality.
 {combined_code}
+Instructions:
+1. Identify severe security/logic issues as 'vulnerabilities' with a 'suggestion' fix.
+2. If no vulnerabilities exist, give 1-2 concise 'improvement_suggestions'.
+3. Assign maintainability and complexity metrics (1-10).
+4. Keep descriptions short to maintain fast execution.
 
-Task:
-1. Detect severe security/logic issues (Vulnerabilities).
-2. If a vulnerability is found, provide a concise 'suggestion' on how to fix it.
-3. If NO vulnerabilities are found in a file, provide a list of 'improvement_suggestions' (clean code, performance, or architecture tips).
-4. Provide metrics for complexity and maintainability (scale 1-10).
-
-Output a JSON array (exactly 1 object per file):
+Return JSON array (1 object per file):
 [
   {{
     "fileName": "exact/path/from/header",
@@ -98,21 +169,19 @@ Output a JSON array (exactly 1 object per file):
         "type": "SQLi/Logic/etc",
         "line": 10,
         "description": "Short explanation",
-        "suggestion": "Specific code fix"
+        "suggestion": "Specific fix"
       }}
     ],
-    "improvement_suggestions": ["Tip 1", "Tip 2"],
+    "improvement_suggestions": ["Tip 1"],
     "metrics": {{"complexity": 3, "maintainability": 8}}
   }}
-]
-"""
-        batch_results = []
+]"""
 
+        batch_results = []
         try:
-            response_text = llm_engine.generate(prompt, max_tokens=8192)
+            response_text = llm_engine.generate(prompt, max_tokens=4096)
             parsed = self._parse_json(response_text)
 
-            # Map results by fileName for easy lookup
             processed_map = {
                 item.get("fileName"): item for item in parsed if isinstance(item, dict)
             }
@@ -120,45 +189,29 @@ Output a JSON array (exactly 1 object per file):
             for fn in file_names:
                 if fn in processed_map:
                     res = processed_map[fn]
-
-                    # Ensure all required keys exist to prevent Android
-                    # serialization errors
                     res.setdefault("vulnerabilities", [])
                     res.setdefault("improvement_suggestions", [])
                     res.setdefault("metrics", {"complexity": 5, "maintainability": 5})
 
-                    # Ensure every vulnerability has a suggestion field
                     for vuln in res["vulnerabilities"]:
                         if "suggestion" not in vuln:
-                            vuln["suggestion"] = (
-                                "Review the implementation logic for improved safety."
-                            )
-
+                            vuln["suggestion"] = "Review implementation logic for safety."
                     batch_results.append(res)
                 else:
-                    # Fallback if the AI skipped a file in its response
                     batch_results.append(
                         {
                             "fileName": fn,
                             "vulnerabilities": [],
-                            "improvement_suggestions": [
-                                "No immediate improvements identified."
-                            ],
+                            "improvement_suggestions": ["No immediate improvements identified."],
                             "metrics": {"complexity": 1, "maintainability": 10},
                         }
                     )
-
         except Exception as e:
             logger.error(f"Batch processing error: {e}")
             error_str = str(e)
             is_config_error = any(
                 x in error_str
-                for x in [
-                    "QUOTA_EXCEEDED",
-                    "RATE_LIMIT_EXCEEDED",
-                    "MISSING_API_KEY",
-                    "GENERATION_FAILED",
-                ]
+                for x in ["QUOTA_EXCEEDED", "RATE_LIMIT_EXCEEDED", "MISSING_API_KEY", "GENERATION_FAILED"]
             )
             for fn in file_names:
                 batch_results.append(
@@ -170,72 +223,44 @@ Output a JSON array (exactly 1 object per file):
                         "error": error_str if is_config_error else "analysis_failed",
                     }
                 )
-
         return batch_results
 
     def _minify_code(self, code: str, filename: str) -> str:
-        """
-        Removes comments, empty lines, and logs to optimize token usage.
-        """
         lines = code.split("\n")
         cleaned_lines = []
-
         is_python = filename.endswith(".py")
-        is_js_style = filename.endswith(
-            (".js", ".ts", ".jsx", ".tsx", ".java", ".kt", ".c", ".cpp")
-        )
+        is_js_style = filename.endswith((".js", ".ts", ".jsx", ".tsx", ".java", ".kt", ".c", ".cpp"))
 
         for line in lines:
             stripped = line.strip()
-
             if not stripped:
                 continue
-
-            # Skip comments
             if is_python and stripped.startswith("#"):
                 continue
             if is_js_style and stripped.startswith(("//", "/*", "*")):
                 continue
-
-            # Skip common logging
-            if any(
-                log in stripped
-                for log in ["console.log", "print(", "logger.", "Log.d", "Log.e"]
-            ):
+            if any(log in stripped for log in ["console.log", "print(", "logger.", "Log.d", "Log.e"]):
                 continue
-
-            # Strip inline comments
             if is_python and " #" in line:
                 line = line.split(" #", 1)[0]
             if is_js_style and " //" in line:
                 line = line.split(" //", 1)[0]
-
-            if not line.strip():
-                continue
-
-            cleaned_lines.append(line.rstrip())
+            if line.strip():
+                cleaned_lines.append(line.rstrip())
 
         return "\n".join(cleaned_lines)
 
     def _parse_json(self, text: str):
-        """
-        Cleans and parses the LLM response into a Python list/dict.
-        """
         text = text.strip()
         if not text:
             return []
-
-        # Clean markdown formatting if present
         text = re.sub(r"^[^[]*\[", "[", text)
         text = re.sub(r"\][^]]*$", "]", text)
         text = text.replace("```json", "").replace("```", "").strip()
-
         try:
             return json.loads(text)
         except json.JSONDecodeError as e:
             logger.warning(f"JSON Decode Error: {e}")
             return []
 
-
-# Instantiate the service
 service = AIReviewerService()