Update app/main.py

app/main.py

@@ -1,233 +1,212 @@
-import os
-import re
+import json
 import logging
-import traceback
-import time
-import asyncio
-from typing import List, Optional, Dict
-from concurrent.futures import ThreadPoolExecutor
-from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, status
-from pydantic import BaseModel
-import uvicorn
-
-load_dotenv()
-
-from app.predictor import classifier, guide_generator, reviewer
+import re
+import hashlib
+import sqlite3
+import threading
+from typing import List
 from app.core.model_loader import llm_engine
 
-logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-
-app = FastAPI(title="GitGud AI Service")
-
-REPO_CACHE: Dict[str, Dict[str, List[float]]] = {}
-executor = ThreadPoolExecutor(max_workers=8)
-
-class FileRequest(BaseModel):
-    fileName: str
-    content: Optional[str] = None
-    repoName: Optional[str] = None
-
-class BatchReviewRequest(BaseModel):
-    files: List[FileRequest]
-
-class GuideRequest(BaseModel):
-    repoName: str
-    filePaths: List[str]
-
-class SearchRequest(BaseModel):
-    query: str
-    embeddings: Optional[Dict[str, List[float]]] = None
-    repoName: Optional[str] = None
-
-class ChatRequest(BaseModel):
-    query: str
-    context: List[Dict[str, str]]
-    repoName: str
-
-def calculate_repo_health(total_vulns: int, avg_maint: float) -> int:
-    base_score = avg_maint * 10
-    penalty = total_vulns * 8
-    return int(max(10, min(100, base_score - penalty)))
-
-def sync_review_worker(file_list: List[FileRequest]):
-    logger.info(f"--- [DEBUG] Processing {len(file_list)} files for code review ---")
-    try:
-        results = reviewer.service.review_batch_code(file_list)
+DB_PATH = "review_cache.db"
+
+class AIReviewerService:
+    def __init__(self):
+        self.batch_size = 3  # Reduced from 10 to speed up individual prompts
+        self.db_lock = threading.Lock()
+        self._init_db()
+
+    def _init_db(self):
+        """Initializes SQLite cache table with proper timeout."""
+        with self.db_lock:
+            with sqlite3.connect(DB_PATH, timeout=30.0) 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:
+        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:
+        with self.db_lock:
+            with sqlite3.connect(DB_PATH, timeout=30.0) 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]):
+        if not hash_result_pairs:
+            return
+        with self.db_lock:
+            with sqlite3.connect(DB_PATH, timeout=30.0) 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:
+        if not files:
+            return []
+
+        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
+
+        # 2. Sequential batch processing (prevents LLM thread lock contention)
+        batches = [
+            uncached_files[i : i + self.batch_size]
+            for i in range(0, len(uncached_files), self.batch_size)
+        ]
+        
+        new_results = []
+        for batch in batches:
+            try:
+                batch_results = self._process_batch(batch)
+                new_results.extend(batch_results)
+            except Exception as e:
+                logger.error(f"Batch processing error: {e}")
+                for f_item in batch:
+                    new_results.append({
+                        "fileName": f_item.fileName,
+                        "vulnerabilities": [],
+                        "improvement_suggestions": [],
+                        "metrics": {"complexity": 0, "maintainability": 0},
+                        "error": "analysis_failed",
+                    })
+
+        new_results_map = {item.get("fileName"): item for item in new_results}
+        to_cache = []
+
+        # 3. Assemble final response and update cache
+        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
-    except Exception as e:
-        logger.error(f"--- [ERROR] Exception during review processing: {e} ---", exc_info=True)
-        raise e
-
-def parse_tree_to_list(raw_tree: str):
-    nodes = []
-    for line in raw_tree.strip().split('\n'):
-        level = line.count('|') + (line.count(' ') // 2)
-        name = re.sub(r'[|└├─]', '', line).strip()
-        if name:
-            nodes.append({
-                "name": name,
-                "type": "file" if '.' in name else "folder",
-                "level": level
-            })
-    return nodes
-
-@app.get("/")
-def health_check():
-    return {
-        "status": "online",
-        "model": "microsoft/codebert-base",
-        "device": getattr(classifier, "device", "cpu"),
-        "cached_repos": list(REPO_CACHE.keys()),
-    }
-
-@app.get("/usage")
-def get_usage():
-    return llm_engine.get_usage_stats()
-
-@app.post("/classify")
-async def classify_file(request: FileRequest):
-    try:
-        result = classifier.predict(request.fileName, request.content)
-        if request.repoName:
-            if request.repoName not in REPO_CACHE:
-                REPO_CACHE[request.repoName] = {}
-            REPO_CACHE[request.repoName][request.fileName] = result["embedding"]
-        return {
-            "fileName": request.fileName,
-            "layer": result["label"],
-            "confidence": result["confidence"],
-            "embedding": result["embedding"]
-        }
-    except Exception as e:
-        logger.error(f"Classify failed: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/review-batch-code")
-async def review_batch_code(request: BatchReviewRequest):
-    try:
-        loop = asyncio.get_running_loop()
-        raw_results = await loop.run_in_executor(executor, sync_review_worker, request.files)
-        final_reviews = raw_results if isinstance(raw_results, list) else [raw_results]
-        return {"results": final_reviews}
-    except Exception as e:
-        logger.error(f"Batch review critical failure: {traceback.format_exc()}")
-        raise HTTPException(status_code=500, detail="Internal processing error")
-
-@app.post("/repo-dashboard-stats")
-async def get_dashboard_stats(request: BatchReviewRequest):
-    try:
-        loop = asyncio.get_running_loop()
-        raw_reviews = await loop.run_in_executor(executor, sync_review_worker, request.files)
-        if not isinstance(raw_reviews, list):
-            raw_reviews = [raw_reviews]
-
-        total_vulns = 0
-        maint_scores = []
-        found_apis = set()
-        api_regex = re.compile(r'(?:get|post|put|delete|patch)\([\'"]\/(.*?)[\'"]', re.IGNORECASE)
-
-        for i, current_review in enumerate(raw_reviews):
-            vulns = current_review.get("vulnerabilities", [])
-            total_vulns += len(vulns)
-            m_score = current_review.get("metrics", {}).get("maintainability", 8.0)
-            maint_scores.append(m_score)
-
-            content = request.files[i].content if i < len(request.files) else None
-            if content:
-                matches = api_regex.findall(content)
-                for match in matches:
-                    found_apis.add(f"/{match}")
-
-        num_files = len(maint_scores)
-        avg_maint = (sum(maint_scores) / num_files) if num_files > 0 else 0
-        health_score = calculate_repo_health(total_vulns, avg_maint)
-
-        return {
-            "repo_health": health_score,
-            "health_label": "Excellent" if health_score > 85 else "Good" if health_score > 60 else "Critical",
-            "security_issues": total_vulns,
-            "performance_ratio": f"{int(avg_maint * 10)}%",
-            "exposed_apis": list(found_apis),
-            "total_files_processed": num_files,
-            "average_maintainability": round(avg_maint, 1)
-        }
-    except Exception as e:
-        logger.error(f"Stats failed: {e}")
-        raise HTTPException(status_code=500, detail="Failed to aggregate metrics")
-
-@app.post("/analyze-file")
-async def analyze_file(request: FileRequest):
-    try:
-        result = classifier.predict(request.fileName, request.content)
-        summary = classifier.generate_file_summary(request.content, request.fileName)
-        tags = classifier.extract_tags(request.content, request.fileName)
-        return {
-            "fileName": request.fileName,
-            "layer": result["label"],
-            "summary": summary,
-            "tags": tags,
-            "embedding": result["embedding"],
-        }
-    except Exception as e:
-        if "429" in str(e):
-            raise HTTPException(status_code=429, detail="Limit Reached")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/semantic-search")
-async def semantic_search(request: SearchRequest):
-    try:
-        embeddings = request.embeddings
-        if not embeddings and request.repoName and request.repoName in REPO_CACHE:
-            embeddings = REPO_CACHE[request.repoName]
-
-        if not embeddings:
-            return {"results": []}
-
-        results = classifier.semantic_search(request.query, embeddings)
-        return {"results": results}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/chat")
-async def chat(request: ChatRequest):
-    try:
-        context_str = ""
-        for item in request.context:
-            context_str += f"--- FILE: {item['fileName']} ---\n{item['content']}\n\n"
-
-        prompt = f"""You are "GitGud AI", an expert software architect.
-Repository: "{request.repoName}"
-CONTEXT: {context_str if request.context else "(NO CODE PROVIDED)"}
-USER QUESTION: {request.query}"""
-
-        response = llm_engine.generate_text(prompt)
-        return {"response": response, "status": "success"}
-    except Exception as e:
-        if "429" in str(e):
-            return {"response": "⚠️ Daily limit reached. Try again in a bit!", "status": "quota_error"}
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/generate-guide")
-async def generate_guide(request: GuideRequest):
-    try:
-        markdown = guide_generator.generate_markdown(request.repoName, request.filePaths)
-        tree_match = re.search(r"Project Structure\n\n(.*?)(?=\n\n|$)", markdown, re.S)
-        structured_tree = []
-        if tree_match:
-            structured_tree = parse_tree_to_list(tree_match.group(1))
-
-        return {
-            "markdown": markdown,
-            "structured_tree": structured_tree,
-            "project_name": request.repoName
-        }
-    except Exception as e:
-        if "429" in str(e):
-            raise HTTPException(status_code=429, detail="AI Quota Exceeded")
-        raise HTTPException(status_code=500, detail=str(e))
-
-if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+
+    def _process_batch(self, batch: list) -> list:
+        combined_code = ""
+        file_names = []
+        for f in batch:
+            raw_content = f.content or ""
+            minified_content = self._minify_code(raw_content[:2000], f.fileName)  # Reduced preview limit to 2000 chars
+            combined_code += f"\n--- FILE: {f.fileName} ---\n{minified_content}\n"
+            file_names.append(f.fileName)
+
+        prompt = f"""Analyze {len(batch)} files. Be ultra-concise.
+{combined_code}
+Return ONLY valid JSON array:
+[
+  {{
+    "fileName": "name",
+    "vulnerabilities": [{{"type": "type", "line": 1, "description": "short", "suggestion": "fix"}}],
+    "improvement_suggestions": ["tip"],
+    "metrics": {{"complexity": 3, "maintainability": 8}}
+  }}
+]"""
+        batch_results = []
+        try:
+            # Capped token length to 1024 max per batch turn
+            response_text = llm_engine.generate(prompt, max_tokens=1024)
+            parsed = self._parse_json(response_text)
+            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]
+                    res.setdefault("vulnerabilities", [])
+                    res.setdefault("improvement_suggestions", [])
+                    res.setdefault("metrics", {"complexity": 5, "maintainability": 5})
+                    batch_results.append(res)
+                else:
+                    batch_results.append({
+                        "fileName": fn,
+                        "vulnerabilities": [],
+                        "improvement_suggestions": [],
+                        "metrics": {"complexity": 1, "maintainability": 10},
+                    })
+        except Exception as e:
+            logger.error(f"Batch processing error: {e}")
+            for fn in file_names:
+                batch_results.append({
+                    "fileName": fn,
+                    "vulnerabilities": [],
+                    "improvement_suggestions": [],
+                    "metrics": {"complexity": 0, "maintainability": 0},
+                    "error": str(e),
+                })
+        return batch_results
+
+    def _minify_code(self, code: str, filename: str) -> str:
+        lines = code.split("\n")
+        cleaned_lines = []
+        is_python = filename.endswith(".py")
+        is_js_style = filename.endswith((".js", ".ts", ".jsx", ".tsx"))
+
+        for line in lines:
+            stripped = line.strip()
+            if not stripped:
+                continue
+            if is_python and stripped.startswith("#"):
+                continue
+            if is_js_style and stripped.startswith(("//", "/*", "*")):
+                continue
+            cleaned_lines.append(stripped)
+
+        return "\n".join(cleaned_lines[:50])  # Cap at 50 lines max per file
+
+    def _parse_json(self, text: str):
+        text = text.strip()
+        if not text:
+            return []
+        text = re.sub(r"^[^\[]*\[", "[", text)
+        text = re.sub(r"\][^\]]*$", "]", text)
+        text = text.replace("```json", "").replace("```", "").strip()
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError:
+            return []
+
+service = AIReviewerService()