Update app/services/reviewer_service.py

app/services/reviewer_service.py

@@ -2,36 +2,84 @@ import json
 import logging
 import re
 from typing import List, Dict
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from app.core.model_loader import llm_engine
 
 logger = logging.getLogger(__name__)
 
+
 class AIReviewerService:
     def __init__(self):
-        pass
+        # 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 = 3
+        self.batch_size = 5
 
     def review_batch_code(self, files: list) -> list:
-        results = []
-        
-        # Process in batches of 5 to stay within Gemini Free Tier limits (15 RPM)
-        batch_size = 5
-        
-        for i in range(0, len(files), batch_size):
-            batch = files[i : i + batch_size]
-            combined_code = ""
-            file_names = []
-            
-            for f in batch:
-                # Minify code to save tokens
-                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)
-                
-                combined_code += f"\n--- FILE: {f.fileName} ---\n{minified_content}\n"
-                file_names.append(f.fileName)
-
-            # Updated prompt to explicitly request suggestions and improvements
-            prompt = f"""
+        """
+        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:
+            return []
+
+        results: list = []
+
+        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)
+                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(
+                            {
+                                "fileName": f.fileName,
+                                "vulnerabilities": [],
+                                "improvement_suggestions": [],
+                                "metrics": {"complexity": 0, "maintainability": 0},
+                                "error": "analysis_failed",
+                            }
+                        )
+
+        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)
+
+            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.
 {combined_code}
 
@@ -50,7 +98,7 @@ Output a JSON array (exactly 1 object per file):
         "type": "SQLi/Logic/etc",
         "line": 10,
         "description": "Short explanation",
-        "suggestion": "Specific code fix" 
+        "suggestion": "Specific code fix"
       }}
     ],
     "improvement_suggestions": ["Tip 1", "Tip 2"],
@@ -58,87 +106,116 @@ Output a JSON array (exactly 1 object per file):
   }}
 ]
 """
-            try:
-                # 8k output tokens for the batch analysis
-                response_text = llm_engine.generate(prompt, max_tokens=8192)
-                batch_results = self._parse_json(response_text)
-                
-                # Map results by fileName for easy lookup
-                processed_map = {item.get('fileName'): item for item in batch_results if isinstance(item, dict)}
-                
-                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."
-                        
-                        results.append(res)
-                    else:
-                        # Fallback if the AI skipped a file in its response
-                        results.append({
-                            "fileName": fn, 
-                            "vulnerabilities": [], 
-                            "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 fn in file_names:
-                    results.append({
-                        "fileName": fn, 
-                        "vulnerabilities": [], 
-                        "improvement_suggestions": [], 
+        batch_results = []
+
+        try:
+            response_text = llm_engine.generate(prompt, max_tokens=8192)
+            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)
+            }
+
+            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."
+                            )
+
+                    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."
+                            ],
+                            "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 fn in file_names:
+                batch_results.append(
+                    {
+                        "fileName": fn,
+                        "vulnerabilities": [],
+                        "improvement_suggestions": [],
                         "metrics": {"complexity": 0, "maintainability": 0},
-                        "error": error_str if is_config_error else "analysis_failed"
-                    })
+                        "error": error_str if is_config_error else "analysis_failed",
+                    }
+                )
 
-        return results
+        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')
+        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_python = filename.endswith(".py")
+        is_js_style = filename.endswith(
+            (".js", ".ts", ".jsx", ".tsx", ".java", ".kt", ".c", ".cpp")
+        )
+
         for line in lines:
             stripped = line.strip()
-            
-            if not stripped: continue
-            
+
+            if not stripped:
+                continue
+
             # Skip comments
-            if is_python and stripped.startswith('#'): continue
-            if is_js_style and stripped.startswith(('//', '/*', '*')): continue
-            
+            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
+            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())
-            
-        return '\n'.join(cleaned_lines)
+
+        return "\n".join(cleaned_lines)
 
     def _parse_json(self, text: str):
         """
@@ -147,17 +224,18 @@ Output a JSON array (exactly 1 object per file):
         text = text.strip()
         if not text:
             return []
-            
+
         # Clean markdown formatting if present
-        text = re.sub(r"^[^[]*\[", "[", text) 
-        text = re.sub(r"\][^]]*$", "]", text) 
+        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()